Skip to content

Simplify dropdims(::Transpose) and insertdims(::PermutedDimsArray) etc. #55381

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.

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

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

Conversation

mcabbott
Copy link
Contributor

@mcabbott mcabbott commented Aug 5, 2024

Before, these wrap the wrapper in a ReshapedArray:

julia> dropdims([1,2,3]'; dims=1)
3-element reshape(adjoint(::Vector{Int64}), 3) with eltype Int64:
 1
 2
 3

julia> insertdims(PermutedDimsArray([1 2; 3 4], (2,1)); dims=2)
2×1×2 reshape(PermutedDimsArray(::Matrix{Int64}, (2, 1)), 2, 1, 2) with eltype Int64:
[:, :, 1] =
 1
 2

[:, :, 2] =
 3
 4

After:

julia> dropdims([1,2,3]'; dims=1)
3-element Vector{Int64}:
 1
 2
 3

julia> insertdims(PermutedDimsArray([1 2; 3 4], (2,1)); dims=2)
2×1×2 PermutedDimsArray(::Array{Int64, 3}, (2, 3, 1)) with eltype Int64:
[:, :, 1] =
 1
 2

[:, :, 2] =
 3
 4

This is not always faster, alone, but the hope is that the simpler return type will usually be faster downstream. (Especially with e.g. CuArrays, where wrapping twice causes them to miss specialised methods.)

julia> P13 = PermutedDimsArray(randn(2,1,3,1,4), (4,5,2,1,3));

julia> @btime dropdims($P13; dims=(1,3));
  189.628 ns (10 allocations: 192 bytes)  # before, ReshapedArray(PermutedDimsArray(Array{Float64, 5}))
  272.895 ns (10 allocations: 256 bytes)  # after, PermutedDimsArray(Array{Float64, 3})

julia> A13 = collect(P13);

julia> @btime dropdims($A13; dims=(1,3));
  172.333 ns (11 allocations: 240 bytes)

It appears that the compiler can predict the type when dims is constant. More stringent tests (or ways to do better) would be welcome:

julia> drop1(x) = dropdims(x; dims=(1,));

julia> @code_warntype drop1(P13)  # before
MethodInstance for drop1(::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}})
  from drop1(x) @ Main REPL[12]:1
Arguments
  #self#::Core.Const(drop1)
  x::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}}
Body::Base.ReshapedArray{Float64, 4, PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}}, NTuple{4, Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}
1%1 = (:dims,)::Core.Const((:dims,))
...

julia> @code_warntype drop1(P13)  # after
MethodInstance for drop1(::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}})
  from drop1(x) @ Main REPL[240]:1
Arguments
  #self#::Core.Const(drop1)
  x::PermutedDimsArray{Float64, 5, (4, 5, 2, 1, 3), (4, 3, 5, 1, 2), Array{Float64, 5}}
Body::PermutedDimsArray{Float64, 4, (4, 2, 1, 3), (3, 2, 4, 1), Array{Float64, 4}}
1%1 = (:dims,)::Core.Const((:dims,))
...

@mcabbott mcabbott added the arrays [a, r, r, a, y, s] label Aug 5, 2024
@mcabbott mcabbott requested a review from nsajko August 23, 2024 17:08
@mcabbott mcabbott requested a review from mbauman August 23, 2024 17:10
@mcabbott
Copy link
Contributor Author

Xref #45793 which added insertdims. And I requested review from a few of those who commented on that PR.

@mcabbott
Copy link
Contributor Author

mcabbott commented Jan 6, 2025

Pre-0.12 bump?

After #56637 this would need to be split into two, but I won't bother if nobody will review it.

@mcabbott mcabbott added the linear algebra Linear algebra label Jan 6, 2025
@DilumAluthge
Copy link
Member

We have moved the LinearAlgebra stdlib to an external repo: https://github.com/JuliaLang/LinearAlgebra.jl

@mcabbott If you think that this PR is still relevant, please open a new PR on the LinearAlgebra.jl repo.

@mcabbott
Copy link
Contributor Author

No, only half of this PR is LinearAlgebra, the other half is Base. It can stay open for discussion?

@mcabbott mcabbott reopened this Jan 13, 2025
@DilumAluthge
Copy link
Member

Yes, makes sense to re-open this PR for the parts that touch Base.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
arrays [a, r, r, a, y, s] linear algebra Linear algebra
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants