Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

MutableTangent #105

Closed
oxinabox opened this issue Jan 19, 2020 · 10 comments · Fixed by #626
Closed

MutableTangent #105

oxinabox opened this issue Jan 19, 2020 · 10 comments · Fixed by #626
Labels
enhancement New feature or request forward-mode Related to use of ChainRules for ForwardMode AD mutability For issues relating to supporting mutability Structural Tangent Related to the `Tangent` type for structured (composite) values

Comments

@oxinabox
Copy link
Member

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.

@oxinabox oxinabox added enhancement New feature or request Structural Tangent Related to the `Tangent` type for structured (composite) values mutability For issues relating to supporting mutability forward-mode Related to use of ChainRules for ForwardMode AD labels Jan 19, 2020
@nickrobinson251
Copy link
Contributor

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
Copy link
Member Author

What's the difference to Composite exactly?

Primary difference is setproperty! will work.

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

oxinabox added a commit that referenced this issue Jan 22, 2020
* PR preview docs

* add link to remind to clean up this code once less evil way is possible

* Update docs/make.jl

* fix typo

* .

* .

* fix docs badge links

* remove debug println

* =highlight readme
@oxinabox
Copy link
Member Author

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

@oxinabox oxinabox changed the title MutableComposite MutableTangent Jul 25, 2023
@oxinabox
Copy link
Member Author

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

@ToucheSir
Copy link
Contributor

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
Copy link
Member Author

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
Copy link
Contributor

ToucheSir commented Jul 27, 2023

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
Copy link
Member Author

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
Copy link
Contributor

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
Copy link
Member Author

oxinabox commented Jul 28, 2023

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request forward-mode Related to use of ChainRules for ForwardMode AD mutability For issues relating to supporting mutability Structural Tangent Related to the `Tangent` type for structured (composite) values
Projects
None yet
Development

Successfully merging a pull request may close this issue.

3 participants