Choosing between `Active` and `Duplicated`

[gdalle]

To compute gradients with Enzyme on a non-mutable array (like SArray), I need to annotate it as Active. Otherwise, if I do this

using StaticArrays, Enzyme
s = SVector(1.0, 2.0)
ds = zero(s)
Enzyme.autodiff(Reverse, sum, Active, Enzyme.Duplicated(s, ds))

there is no way to recover the derivative because ds doesn't change.

My default setting in DI is to mark arrays as Duplicated. Is there a user-facing function in Enzyme that I could call to pick a better-suited annotation?

Ping @ExpandingMan

[ExpandingMan]

For whatever reason this seems to be a much bigger problem for reverse mode. Forward mode usually seems to "just work" and I don't always understand why. I think more extensive documentation about all the various mode arguments would be useful, as well as more examples.

[wsmoses]

yeah we should probably add better docs on, offhand probably see https://enzyme.mit.edu/getting_started/CallingConvention/#types and https://enzyme.mit.edu/index.fcgi/julia/stable/faq/#Mixed-activity

Forward mode doesn't have active so there's never a choice needed

[gdalle]

The typical problem is for the generic pullback inside DI. In this line, I pick a return activity according to a very simple rule of thumb: Number gets Active, anything else is Duplicated.

https://github.com/gdalle/DifferentiationInterface.jl/blob/3470e1424d667ffe79112c3b8429c9c0ca0e81e5/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl#L74

But it is wrong for SArrays, and for mixed activity. I presume Enzyme has some automatic activity guesser inside?

[wsmoses]

Yup, an internal function Enzyme.Compiler.active_reg_inner. But note that it is definitionally type unstable (though it tries its best to be)

[gdalle]

Would you be open to documenting it so that it can be used by DI? Type instability is less of an issue thanks to the preparation mechanism, which we only use at the beginning.

[wsmoses]

sure to docs, but it definitely should remain an internal function since it takes a recursive descent state that might need to be modified as we build out more activity analysis infra and I don't want to force a breaking change for that.

[gdalle]

If it is used by DI it will have to be part of the public API. But just being part of the public API doesn't necessarily mean that its behavior has to stay exactly the same. I would be fine with a public function whose docstring says something like

    Enzyme.Compiler.active_reg_inner

Pick the most appropriate activity annotation for a given object and AD mode.
This function is subject to change as Enzyme's activity analysis progresses, but its goal will remain the same.

[wsmoses]

I mean if you're willing to keep up with any changes to the ABI in patch releases (the threshold which I would refer to something as an internal function), I'm happy to put a docstring on it.

I don't expect that to happen often (or frankly much at all), but it's definitely not something that should force an Enzyme breaking change since as we just saw that often requires a whole ecosystem sweep.

[gdalle]

What would be the consequences of such ABI changes on the activity analysis? I still don't fully understand what the ABI does ^^ #1345

[danielwe]

There's the higher-level Enzyme.guess_activity(::Type, ::Enzyme.Mode) that seems like could be made public without much trouble? https://github.com/EnzymeAD/Enzyme.jl/blob/4d4c546dbafb048f7fe73195925bbd81a8d1105a/src/compiler.jl#L859-L860

[wsmoses]

Oh right I forgot that existed, yeah that one is fine to make public, since we can still change any of the args to active reg inner with ease

[gdalle]

Opened #1955 with a teeny tiny docstring