feat: stacked batchduplicated

[avik-pal]

This is mostly for use in reactant

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 12 lines in your changes missing coverage. Please review.

Project coverage is 75.22%. Comparing base (037dfed) to head (51453ef).
Report is 401 commits behind head on main.

Files with missing lines	Patch %	Lines
lib/EnzymeCore/src/EnzymeCore.jl	0.00%	12 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2418      +/-   ##
==========================================
+ Coverage   67.50%   75.22%   +7.72%     
==========================================
  Files          31       56      +25     
  Lines       12668    16939    +4271     
==========================================
+ Hits         8552    12743    +4191     
- Misses       4116     4196      +80     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[wsmoses]

so one caveat here, can you grep through the code for places with BatchDuplicated and check that we add the requisite condition for StackedBatchDuplicated.

Even more ideally we would have some interface function is_batched or something that can be used by either

[avik-pal]

@inline remove_innerty(::Type{<:BatchDuplicated}) = Duplicated
@inline remove_innerty(::Type{<:BatchDuplicatedNoNeed}) = DuplicatedNoNeed

what are these for?

[wsmoses]

Active{Float64} -> Active

[vchuravy]

If it isn't supported by Enzyme, should it then be defined here? I.e. is there a benefit for this to be in EnzymeCore and could it just live in Reactant?

[wsmoses]

I mean honestly this would be nice to support from Enzyme proper as well from a UX perspecive (and could make things like jacobian easier since you give it the full matrix)

[wsmoses]

we'd probably need to do a similar shim layer like we do for mixedduplicated creating subpieces from the bigger one/etc

[vchuravy]

On Julia 1.11 this is possible due to the Memory Change.

function tuple_of_vectors(M::Matrix{T}, shape) where {T}
        n, m = size(M)
        return ntuple(m) do i
            vec = Base.wrap(Array, memoryref(M.ref, (i - 1) * n + 1), (n,))
            reshape(vec, shape)
        end
    end

    function mul!(Out::AbstractMatrix, J::JacobianOperator, V::AbstractMatrix)
        @assert size(Out, 2) == size(V, 2)
        out = tuple_of_vectors(Out, size(J.res))
        v = tuple_of_vectors(V, size(J.u))

        N = length(out)
        autodiff(
            Forward,
            maybe_duplicated(J.f, Val(N)), Const,
            BatchDuplicated(J.res, out),
            BatchDuplicated(J.u, v),
            maybe_duplicated(J.p, Val(N))
        )
        return nothing
    end

[wsmoses]

yeah that would be nice