You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
ChainRulesCore.jl offers experimental support for mutation, targeting use in forward mode AD.
4
+
(Mutation support in reverse mode AD is more complicated and will likely require more changes to the interface)
5
+
6
+
!!! warning "Experimental"
7
+
This page documents an experimental feature.
8
+
Expect breaking changes in minor versions while this remains.
9
+
It is not suitable for general use unless you are prepared to modify how you are using it each minor release.
10
+
It is thus suggested that if you are using it to use _tilde_ bounds on supported minor versions.
11
+
12
+
13
+
## `MutableTangent`
14
+
The [`MutableTangent`](@ref) type is designed to be a partner to the [`Tangent`](@ref) type, with specific support for being mutated in place.
15
+
It is required to be a structural tangent, having one tangent for each field of the primal object.
16
+
17
+
Technically, not all `mutable struct`s need to use `MutableTangent` to represent their tangents.
18
+
Just like not all `struct`s need to use `Tangent`s.
19
+
Common examples away from this are natural tangent types like for arrays.
20
+
However, if one is setting up to use a custom tangent type for this it is sufficiently off the beaten path that we can not provide much guidance.
21
+
22
+
## `zero_tangent`
23
+
24
+
The [`zero_tangent`](@ref) function functions to give you a zero (i.e. additive identity) for any primal value.
25
+
The [`ZeroTangent`](@ref) type also does this.
26
+
The difference is that [`zero_tangent`](@ref) is in general full structural tangent mirroring the structure of the primal.
27
+
To be technical the promise of [`zero_tangent`](@ref) is that it will be a value that supports mutation.
28
+
However, in practice[^1] this is achieved through in a structural tangent
29
+
For mutation support this is important, since it means that there is mutable memory available in the tangent to be mutated when the primal changes.
30
+
To support this you thus need to make sure your zeros are created in various places with [`zero_tangent`](@ref) rather than []`ZeroTangent`](@ref).
31
+
32
+
33
+
34
+
It is also useful for reasons of type stability, since it forces a consistent type (generally a structural tangent) for any given primal type.
35
+
For this reason AD system implementors might chose to use this to create the tangent for all literal values they encounter, mutable or not,
36
+
and to process the output of `frule`s to convert [`ZeroTangent`](@ref) into corresponding [`zero_tangent`](@ref)s.
37
+
38
+
## Writing a frule for a mutating function
39
+
It is relatively straight forward to write a frule for a mutating function.
40
+
There are a few key points to follow:
41
+
- There must be a mutable tangent input for every mutated primal input
42
+
- When the primal value is changed, the corresponding change must be made to its tangent partner
43
+
- When a value is returned, return its partnered tangent.
44
+
- If (and only if) primal values alias, then their tangents must also alias.
45
+
46
+
### Example
47
+
For example, consider the primal function with:
48
+
1. takes two `Ref`s
49
+
2. doubles the first one in place
50
+
3. overwrites the second one's value with the literal 5.0
51
+
4. returns the first one
52
+
53
+
54
+
```julia
55
+
functionfoo!(a::Base.RefValue, b::Base.RefValue)
56
+
a[] *=2
57
+
b[] =5.0
58
+
return a
59
+
end
60
+
```
61
+
62
+
The frule for this would be:
63
+
```julia
64
+
function ChainRulesCore.frule((_, ȧ, ḃ), ::typeof(foo!), a::Base.RefValue, b::Base.RefValue)
65
+
@assert ȧ isa MutableTangent{typeof(a)}
66
+
@assert ḃ isa MutableTangent{typeof(b)}
67
+
68
+
a[] *=2
69
+
ȧ.x *=2# `.x` is the field that lives behind RefValues
70
+
71
+
b[] =5.0
72
+
ḃ.x =zero_tangent(5.0) # or since we know that the zero for a Float64 is zero could write `ḃ.x = 0.0`
73
+
74
+
return a, ȧ
75
+
end
76
+
```
77
+
78
+
Then assuming the AD system does its part to makes sure you are indeed given mutable values to mutate (i.e. those `@assert`ions are true) then all is well and this rule will make mutation correct.
79
+
80
+
[^1]:
81
+
Further, it is hard to achieve this promise of allowing mutation to be supported without returning a structural tangent.
82
+
Except in the special case of where the struct is not mutable and has no nested fields that are mutable.
0 commit comments