rrule for stack

[CarloLucibello]

stack will land in julia v1.9 thanks to JuliaLang/julia#43334

cc @mcabbott

[mcabbott]

What I think this misses is that multi-dim containers are allowed, and dims=: their dimensions -- it's not the same as any one integer dims:

julia> stack([(1,2,3) (4,5,6); (7,8,9) (10,11,12)]) |> summary
"3×2×2 Array{Int64, 3}"

julia> stack([(1,2,3) (4,5,6); (7,8,9) (10,11,12)]; dims=2) |> summary
"3×4 Matrix{Int64}"

julia> stack([(1,2,3) (4,5,6); (7,8,9) (10,11,12)]; dims=1) |> summary
"4×3 Matrix{Int64}"

What I sketched on while fidding with #671 was this, but not sure I tested it carefully. Using eachslice ought to mean that second derivatives work.

function frule((_, ẋ), ::typeof(stack), x; dims = :)
    return stack(x; dims), stack(ẋ; dims)
end

# Other iterable X also allowed, maybe this should be wider?
function rrule(::typeof(stack), X::AbstractArray; dims = :)
    Y = stack(X; dims)
    sdims = if dims isa Colon
        N = ndims(Y) - ndims(X)
        X isa AbstractVector ? ndims(Y) : ntuple(i -> i + N, ndims(X))
    else
        dims
    end
    project = ProjectTo(x)
    function stack_pullback(dY)
        dX = collect(eachslice(unthunk(dY); dims = sdims))
        return (NoTangent(), project(dX))
    end
    return Y, stack_pullback
end

# This wants #671, but ought to work with Zygote already?
function rrule(config, ::typeof(stack), f, args...; dims = :)
    y, unmap = rrule_via_ad(config, map, f, args...)
    z, unstack = rrule(stack, y)
    function stack_pullback_f(dz)
        _, dy = unstack(dz)
        _, df, dargs... = unmap(dy)
        return (NoTangent(), df, dargs...)
    end
    return z, stack_pullback_f
end

I guess multiple container dimensions here needs the new eachslice, also 1.9 but not in Compat yet, maybe that's OK. I believe that returning views is also OK, and probably all the cat rules should do this.

Handling generic containers prob. needs something like Base._iterator_axes:

_ndims(x) = _ndims(x, IteratorSize(x))
_ndims(x, ::HasLength) = 1
_ndims(x, ::HasShape{N}) where N = N

[mcabbott]

Something fails on CI but not locally. The testing setup hides the error, and doesn't print a useful line number:

stack: Error During Test at /home/runner/work/ChainRules.jl/ChainRules.jl/test/rulesets/Base/array.jl:420
[156](https://github.com/JuliaDiff/ChainRules.jl/actions/runs/3301482615/jobs/5446999788#step:6:159)
  Got exception outside of a @test
[157](https://github.com/JuliaDiff/ChainRules.jl/actions/runs/3301482615/jobs/5446999788#step:6:160)

[CarloLucibello]

Something fails on CI but not locally.

thanks Michael, there was a missing import in the tests. I think this is ready to go if tests pass.

[CarloLucibello]

failures not related

[CarloLucibello]

@mcabbott is there something blocking this?

[CarloLucibello]

bump

[mcabbott]

This looks fine, I think a few things can be compacted.