Skip to content

Commit

Permalink
fix column- to row-major layout conversion on Reactant.promote_to call
Browse files Browse the repository at this point in the history
  • Loading branch information
mofeing committed Nov 26, 2024
1 parent de4748a commit 08aeb52
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion ext/TenetReactantExt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -201,7 +201,14 @@ end

Tenet.contract(a::Tensor, b::Tensor{T,N,TracedRArray{T,N}}; kwargs...) where {T,N} = contract(b, a; kwargs...)
function Tenet.contract(a::Tensor{Ta,Na,TracedRArray{Ta,Na}}, b::Tensor{Tb,Nb}; kwargs...) where {Ta,Na,Tb,Nb}
return contract(a, Tensor(Reactant.promote_to(TracedRArray{Tb,Nb}, parent(b)), inds(b)); kwargs...)
# TODO change to `Ops.constant` when Ops PR lands in Reactant
# apparently `promote_to` doesn't do the transpostion for converting from column-major (Julia) to row-major layout (MLIR)
# currently, we call permutedims manually
return contract(
a,
Tensor(Reactant.promote_to(TracedRArray{Tb,Nb}, permutedims(parent(b), collect(Nb:-1:1))), inds(b));
kwargs...,
)
end

end

0 comments on commit 08aeb52

Please sign in to comment.