Documentation about the reverse rule customization needs to be improved
- Dominant language
- Julia
- Stars
- 586
- Forks
- 108
- Avg merge
- 1d 5h
- Merged PRs (30d)
- 44
Description
@vchuravy got the following code for differentiation the `einsum!` function in OMEinsum work. He also pointed out that the relevant documentation could be improved. Hope this code snippet helps.
```julia
using Enzyme, Enzyme.EnzymeRules, OMEinsum
function EnzymeRules.augmented_primal(
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof(einsum!)}, ::Type,
code::Const, xs::Duplicated, ys::Duplicated, sx::Const, sy::Const, size_dict::Const)
@assert sx.val == 1 && sy.val == 0 "Only α = 1 and β = 0 is supported, got: $sx, $sy"
# Compute primal
if EnzymeRules.needs_primal(config)
primal = func.val(code.val, xs.val, ys.val, sx.val, sy.val, size_dict.val)
else
primal = nothing
end
# Save x in tape if x will be overwritten
if EnzymeRules.overwritten(config)[3]
tape = copy(xs.val)
else
tape = nothing
end
shadow = ys.dval
return EnzymeRules.AugmentedReturn(primal, shadow, tape)
end
function EnzymeRules.reverse(config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof(einsum!)}, dret::Type{<:Annotation}, tape,
code::Const,
xs::Duplicated, ys::Duplicated, sx::Const, sy::Const, size_dict::Const)
xval = EnzymeRules.overwritten(config)[3] ? tape : xs.val
for i=1:length(xs.val)
xs.dval[i] .+= OMEinsum.einsum_grad(OMEinsum.getixs(code.val),
xval, OMEinsum.getiy(code.val), size_dict.val, conj(ys.dval), i)
end
return (nothing, nothing, nothing, nothing, nothing, nothing)
end
x = randn(3, 3);
y = randn(3);
gx = zero(x);
gy = zero(y);
function testf2(x)
y = zeros(size(x, 1))
einsum!(ein"ii->i", (x,), y, 1, 0, Dict('i'=>3))
return sum(y)
end
autodiff(ReverseWithPrimal, testf2, Duplicated(x, gx))
gx
```
The function signature of `einsum!` is
```julia
einsum!(code::EinCode, xs::Tuple, y, sx, sy, size_dict::Dict=get_size_dict(getixs(code), xs))
```
The input `y` is directly changed, and the return value is the same as `y`.
Contributor guide
Assessment
This issue has not been assessed yet.