align: refine_names + permutedims

[oxinabox]

I have been thinking about an option that i call align
which combined permutedims with refine_named.

With the following behavour:

using Test

function check(a,b)
    @test a == b
    @test dimnames(a) == dimnames(b)
end


@testset "Align" begin
    
    # On fully named it is `permutedims`
    @test check(
        align(NamedDimsArray{(:A, :B)(ones(3,5)), (:B, :A)),
        NamedDimsArray{(:B, :A)}(ones(5,3))
    )

    # on fully unnamed it is `refine_names`
    @test check(
        align(ones(3,5), (:B, :A)),
        NamedDimsArray{(:B, :A)}(ones(3,5))
    )

    # On partially named it is `permutedims` + `refine_names`
    @test check(
        align(NamedDimsArray{(:A, :_)(ones(3,5)), (:B, :A)),
        NamedDimsArray{(:B, :A)}(ones(5,3))
    )

    # On partially named it defaults minimizing movement of unnamed
    # (This one is confusing to think about)
    @test check(
        align(NamedDimsArray{(:A, :_, :_)(ones(1,3,5)), (:C, :B, :A)),
        NamedDimsArray{(:C, :B, :A)}(ones(3,5,1))
    )

    # Similar behavior as `refine_names` if the align names have wildcards
    @test check(
        align(NamedDimsArray{(:A, :B)(ones(3,5)), (:_, :A)),
        NamedDimsArray{(:B, :A)}(ones(5,3))
    )

    # Sometimes this gets complex and hurts me:
    @test check(
        align(NamedDimsArray{(:A, :B, :_)(ones(1,3,5)), (:C, :_, :A)),
        NamedDimsArray{(:C, :B, :A)}(ones(5,3,1))
    )

    # ununionable cases still error:
    @test_throws DimensionMismatch align(NamedDimsArray{(:A, :B)(ones(3,5)), (:NOPE, :A))
end

This would generally replace most uses of refine_names as the thing you call if you are not sure if you are going to have names.

It would mean as long as you followed convention of what each dimension meant when using general AbstractArrays all is well.
But when you use NamedDimsArrays you can be more relaxed and get the order wrong, knowing it will fix it for you.

Note: this is subtly different from PyTorch's align_to
which only supports the permutedims like features, not the refine_names features, and errors if applied to non-named Tensors.

We could call it align_and_refine maybe instead.
but we don't have a align as we basically use permutedims for that.

Once we have this we can also add:

align_as(x, y) = align(x, dimnames(y))

and maybe for matmul:

reverse_align_as(x, y) = align(x, reverse(dimnames(y)))

[mcabbott]

Edit: Isn't what you're describing just a method of permutedims which takes wildcards? How does it differ from #31?

Earlier: I made an align function which works differently, here. The rule is that Aa = align(A, B) does lazy permutations such that Aa .+ B and sum!(B, Aa) will work. It is happy to insert trivial dimensions, which it labels :_. On your examples:

julia> using NamedPlus

julia> align(NamedDimsArray{(:A, :B)}(ones(3,5)), (:B, :A)) |> summary
"5×3 NamedDimsArray(::LinearAlgebra.Transpose{Float64,Array{Float64,2}}, (:B, :A))"

julia> align(NamedDimsArray{(:A, :B)}(ones(3,5)), (:NOPE, :A)) |> summary
"1×3×5 NamedDimsArray(TransmutedDimsArray(::Array{Float64,2}, (0, 1, 2)), (:_, :A, :B))"

# and with wildcards:

julia> align(NamedDimsArray{(:A, :_)}(ones(3,5)), (:B, :A)) |> summary
"1×3×5 NamedDimsArray(TransmutedDimsArray(::Array{Float64,2}, (0, 1, 2)), (:_, :A, :_))"

julia> align(NamedDimsArray{(:A, :_, :_)}(ones(1,3,5)), (:C, :B, :A)) |> summary
ERROR: input must have unique names, got (:A, :_, :_)
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] align(::NamedDimsArray{(:A, :_, :_),Float64,3,Array{Float64,3}}, ::Tuple{Symbol,Symbol,Symbol}) at /Users/me/.julia/dev/NamedPlus/src/permute.jl:27
 [3] top-level scope at REPL[70]:1

I'm not precisely sure how wildcards in A or B should be handled. The above list doesn't quite match the uploaded code, which I think simply forbids wildcards.
I also made repeated names an error, but a smarter implementation could perhaps allow that.

It always returns a view, often a slightly generalised (and faster) variant of PermutedDimsArray.

[oxinabox]

Edit: Isn't what you're describing just a method of permutedims which takes wildcards? How does it differ from #31?

A key difference is the following case:

    # on fully unnamed it is `refine_names`
    @test check(
        align(ones(3,5), (:B, :A)),
        NamedDimsArray{(:B, :A)}(ones(3,5))
    )

On plain arrays it turrns them into NamedDimsArrays

[mcabbott]

OK. But it would be entirely natural that a permutedims which accepted wildcards would regard a plain array as being all-wildcards. (Although that sounds like a bug-magnet... you're asking to permute to match a pattern, but accepting whatever you get.)

[oxinabox]

We can't redefine permutedims on generic arrays.
That would be type piracy.

It's no more a bug magnet then the current way refine_names is used.
See the read me.

[oxinabox]

Anytime an array has more than one wildcard things get confusing.
I recommend avoiding getting in that situation in the first place.
Which mostly one can do since any individual op at most introduces one wildcard.

[mcabbott]

Yes, good point about piracy. But IMO this still sounds like a variant of permutedims. I can see the value of guaranteeing that a matrix X which may have two names or may just have one, and may be transposed or not, ends up permuted to (:a, :b) before your function proceeds. But if if may not have any names at all, then it no longer sounds like auto-permutation, you are just trusting the order it arrived in, and applying names.

It seems to break the rule of this package that (1) operations not mentioning names should do exactly what they do on ordinary arrays, but propagate names, and error if mismatch, and (2) operations whose result depends on the given names only work on named arrays, and don't guess. It would be a huge bug-magnet if sum(A, dims=:a) was happy to just pick the first dimension and call it :a when acting on an ordinary array.

Anyway you will be shocked to hear that I had a branch playing with wildcards at some point... https://github.com/mcabbott/NamedDims.jl/blob/wild/test/name_core.jl#L90 . This tries to allow them in source or destination or both at once, but I think it got messy.