EnzymeAD / EnzymeAD/Enzyme.jl

Documentation about the reverse rule customization needs to be improved

Open
#2,132 4 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.