Multiple active arguments?

[gdalle]

With v0.6, DI supports any number of context arguments (constants or caches), but only a single active argument x (which comes first).
The last step to full generality is support for any number of active arguments. This would improve performance e.g. with Enzyme, and also (apparently) avoid some differentiation errors.

The issue is that for multiple active arguments, every operator output must become a tuple instead of a single object. In addition to being another breaking change, that might slow down backends which do not support multiple arguments.

My workaround idea for now is:

1. force the first argument to always be active
2. if there are no other active arguments (easy to check by constraining the Vararg), return single objects
3. otherwise, return tuples

Related past discussions:

• Multiple arguments / activities? #311
• Rethink argument order #403

[wsmoses]

recopying from slackhole:

using Enzyme, BenchmarkTools, ADTypes, DifferentiationInterface

mul(a, b) = @inbounds (a[1] * b[1])

const a = [2.0]
const b = [3.0]
Enzyme.gradient(Reverse, mul, a, b)

@btime Enzyme.gradient($Reverse, $mul, $a, $b) # 52.496 ns (2 allocations: 128 bytes)

fuse(tup) = mul(tup[1], tup[2])

function bench_di(backend, a, b)
   prep = DifferentiationInterface.prepare_gradient(fuse, backend, (a, b))
   @btime DifferentiationInterface.gradient($fuse, $prep, $backend, $((a, b)))
   return nothing
end

bench_di(AutoEnzyme(), a, b) #  131.136 ns (4 allocations: 464 bytes)

@btime Enzyme.gradient(Reverse, fuse, $((a, b)) #  129.619 ns (4 allocations: 464 bytes)

[wsmoses]

and

using Enzyme, BenchmarkTools, ADTypes, DifferentiationInterface, StaticArrays

mul(a, b) = @inbounds a[1] * sum(b)

const a = [2.0]
const b = @SMatrix ones(8,8)
Enzyme.gradient(Reverse, mul, a, b)

@btime Enzyme.gradient($Reverse, $mul, $a, $b) #  7.183 μs (13 allocations: 3.45 KiB)

fuse(tup) = mul(tup[1], tup[2])

function bench_di(backend, a, b)
   prep = DifferentiationInterface.prepare_gradient(fuse, backend, (a, b))
   @btime DifferentiationInterface.gradient($fuse, $prep, $backend, $((a, b)))
   return nothing
end

bench_di(AutoEnzyme(), a, b) #   27.910 μs (164 allocations: 11.50 KiB)

@btime Enzyme.gradient(Reverse, fuse, $((a, b))) #  27.810 μs (163 allocations: 10.97 KiB)