Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

adding transpose!() to linear algebra #808

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

adding transpose!() to linear algebra #808

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions src/stdlibs/LinearAlgebra.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -397,4 +397,21 @@ function LinearAlgebra._kron!(C::AnyTracedRMatrix, A::AnyTracedRMatrix, B::AnyTr
return C
end

LinearAlgebra.transpose(a::AnyTracedRArray) = error("transpose not defined for $(typeof(a)).")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there a reason to add this? Shouldn't it fall back to the inplace versions?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The stdlib LinearAlgebra.jl has this:
transpose(a::AbstractArray) = error("transpose not defined for $(typeof(a)). Consider using `permutedims` for higher-dimensional arrays.")

The transpose!() docstring also mentions about not supporting the in-place version.

"""
transpose!(dest,src)

Transpose array src and store the result in the preallocated array dest, which should
have a size corresponding to (size(src,2),size(src,1)). No in-place transposition is
supported and unexpected results will happen if src and dest have overlapping memory
regions.
"""

So yeah, I was just trying to mimic it exactly.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i guess that then you can leave it unimplemented and it will call the default method which will error like is supposed to


function LinearAlgebra.transpose!(B::AnyTracedRVector, A::AnyTracedRMatrix)
LinearAlgebra.check_transpose_axes((size(B,1), size(B,2)), size(A))
set_mlir_data!(B, get_mlir_data(Ops.reshape(A, length(B))))
end

function LinearAlgebra.transpose!(B::AnyTracedRMatrix, A::AnyTracedRVector)
LinearAlgebra.check_transpose_axes(size(B), (size(A, 1), size(A, 2)))
set_mlir_data!(B, get_mlir_data(Ops.broadcast_in_dim(A, [2], [1, length(A)])))
end

function LinearAlgebra.transpose!(B::AnyTracedRMatrix, A::AnyTracedRMatrix)
LinearAlgebra.check_transpose_axes(size(B), size(A))
set_mlir_data!(B, get_mlir_data(Ops.transpose(A, [2,1])))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@avik-pal @wsmoses Should we enforce same eltype for both args? if not, there can be undetected problems when we pass two arrays of same size but different eltype.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should do a promote and set here. Base implements it like B[j,i] = f(A[i,j]) so its auto promoted

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

noted. I'll go fix this.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i know we do promote here but is it the correct thing? e.g. in transpose!(::TracedRArray{Int}, ::TracedRArray{Float64}), promotion of eltype will be Float64 and the result we are storing will be wrong.

or should the promotion handle it and error in that case?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think let's update Ops.convert to throw errors in cases where Julia would also throw an error.

Then we just do Ops.convert(..., Ops.transpose(A, [2,1]))

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wrestling with the Ops.convert() for a while now. It looks like the stablehlo spec is that it performs unsafe truncation by default while Julia does not. And I could not find any Julia function that takes T1, T2 and be able to tell whether type conversion will be safe.

Julia seems to determine if the conversion is OK (like throwing InexactError or not) during run-time by explicitly converting the value to type T (i.e. 1.0 -> 1 is ok, for example).

In the case of TracedRNumber, maybe we can just do T(x) directly and let Julia throw error - if any (not sure yet how to go from mlir_dat -> Julia primitives).

In case of TracedRArray, I don't really know. Is there anything or some obvious solutions I am missing?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah I see what's going on. Unfortunately stablehlo doesn't seem to have conversion mode supported (openxla/stablehlo#180). We can do the convert here and just defer to stablehlo to handle the truncation.

cc @wsmoses any thoughts here?

end

end
26 changes: 26 additions & 0 deletions test/integration/linear_algebra.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,3 +183,29 @@ end
end
end
end

@testset "transpose!" begin
v = zeros(5)
M = rand(1, 5)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would you mind not using rand? we had problems tracking numerical stability and bugs due to rand.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No problem. I will make a change to that. Also, the rest of the test set uses rand. Should I also change them?

v_ra = Reactant.to_rarray(v)
M_ra = Reactant.to_rarray(M)

@jit transpose!(v_ra, M_ra)
@test v_ra ≈ transpose!(v, M)

v = rand(7)
M = zeros(1, 7)
v_ra = Reactant.to_rarray(v)
M_ra = Reactant.to_rarray(M)

@jit transpose!(M_ra, v_ra)
@test M_ra ≈ transpose!(M, v)

A = rand(3, 7)
B = rand(7, 3)
A_ra = Reactant.to_rarray(A)
B_ra = Reactant.to_rarray(B)
@jit transpose!(B_ra, A_ra)
@test B_ra ≈ transpose!(B, A)
end

Loading