MutableTangent

[oxinabox]

Mutable primal types need mutable differentials.

This is one of the things needed to support mutability in forwards mode.
(For reverse mode this is not enough).

We should have a MutableComposite type for them.

Whether it is backed by a Dict; or backed by a NamedTuple that gets swapped out everytime setproperty is called, idk.

[nickrobinson251]

What's the difference to Composite exactly?
How's this different from elementwise_add and _update! (which needs to be changed to work with Composite)

[oxinabox]

What's the difference to Composite exactly?

Primary difference is setproperty! will work.

(also accumulate! will work (both directly, and with InplaceThunks))

[oxinabox]

Might also be used for Dict
though slightly different concern to structs since might have hundreds of elements.

[oxinabox]

This would now be called MutableTangent{T} so I have updated the title

[ToucheSir]

With JuliaDiff/Diffractor.jl#189 in mind, how do you think type stability should be addressed? Currently field values could be any of AbstractZero, structural tangents or natural tangents (including natural zeros).

[oxinabox]

It is not an additional concern.
It remains just as much of a concern as the current immutable Tangent{T} is -- we remain in the small union space for unions across multiple methods, and individual methods remain (mostly) type stable.
If you want to talk more about type stability concerns do open another issue.

[ToucheSir]

As a rule writer who would like to know what types to expect/output/enforce when working with MutableTangents, this small union space is exactly what I was curious to hear your thoughts on! In particular, I was hoping to have my mental model of this corrected.

Here are all the different types I understand are allowed as cotangents for structural Tangent{T} field values:

1. NoTangent/ZeroTangent
2. NotImplemented
3. Natural tangents of type T
4. Natural tangents of a type other than T (as noted in the docs)
5. User-defined tangent types <:AbstractZero or <:AbstractTangent
6. Structural tangents (repeat above choices recursively)

AIUI this proliferation of different types is less of a worry for immutable structural tangents because they can't be updated in place, but more important for MutableTangent because it needs to specify a wide enough type for its backing value fields ahead of time to not error when someone calls setproperty!.

I couldn't find a definitive answer on the small union optimization limit (https://discourse.julialang.org/t/performance-tips-for-union-types/13928/4 mentions 3-255 and JuliaLang/julia#44131 demonstrates 5), but I presume allowing user-defined tangent types in there breaks the optimization because there are theoretically infinite subtypes of both <:AbstractZero or <:AbstractTangent? So presumably when one calls setproperty! on a MutableTangent, some kind of normalization would be required to narrow the value down to a smaller set of known types which can benefit from small union optimization? Depending on how low the limit is in practice, I'd imagine even more normalization (e.g. converting NoTangent to ZeroTangent or vice versa for storage, controlling natural tangent types to a single canonical natural tangent type) would be required?

Apologies if this is too early in the design phase to ask. I saw this issue getting some recent activity and assumed you had some more concrete thoughts on it :)

[oxinabox]

Ah I see.
So the answer is for mutable structs, never use ZeroTangent -- because the primal might get mutated, and so you might need to mutate the tangent.
Potentially if the AD can do an analysis to prove it never actually is mutated and never escapes it might then use ZeroTangent, these proofs are possible with existing effect system i think, but for rule authors this doesn't apply.
Always use either a natural tangent (which is also mutable) or MutableTangent.
And almost always applicable rule: use Array if the primal is an Array, otherwise use MutableTangent.

So in general we don't have to worry about type stability as long as each method is type stable.
And that is almost always the case: a given methodinstance always returns a the same type, or maaaybe occationally one of two types if one branch is hard-zero. So we don't run into that limit.
Remember: type stability is per method instance.
In practice, we might need to do something for sake of AD engine, but there is nothing much can do about that pre-2.0. Diffractor though is getting specific optimizations around small unions to make it more robust, but that's another discussion.
This is off the topic of mutation though.

for mutatiom you just need to ensure mutable tangent for mutable primal.

[ToucheSir]

Makes sense, thanks. I guess my question was more around what type constraints should be placed on the backing type of MutableTangent to allow different types to be written to its fields over time while still keeping the small union optimization for those. e.g. if I start with a canonicalized MutableTangent{T} with all fields in the backing set to ZeroTangent, presumably the actual backing NamedTuple would require a value type wider than ZeroTangent so that rules or AD can fill in fields later on? Otherwise, my understanding is that you could not swap out the underlying backing type B from MutableTangent{T, B<:NamedTuple} without leaving the value type of B as Any or some abstract type.

[oxinabox]

So my plan was for it to be backed by a NamedTuple, and to work just like the immutable one.
So empty but using getproperty to return ZeroTangent (though that's an implementation detail).
When you mutate it, replace that NamedTuple.

But yes, like you say yes, swapping out the named tuple like that would require that the field can accept them.
Very good point.
For a first pass I am very tempted to just set it to Any. But maybe we do need to instead have a suitable automatically determined types.
But we can iterate on that once it is in.