FluxML / FluxML/NNlib.jl

Enzyme: scatter! reverse rule only supports + and -, other ops throw MethodError

Open
#765 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
258
Forks
138
Avg merge
2d 1h
Merged PRs (30d)
3

Description

`scatter` and `scatter!` accept `+`, `-`, `*`, `/`, `max`, `min` and `mean`, and gradients work for all of them under Zygote. Under Enzyme only `+` and `-` work — the other five throw a `MethodError` in the reverse pass.

The forward pass succeeds, so the failure only shows up when the gradient is taken.

### Reproducer

```julia
using NNlib, Enzyme, Zygote
using Statistics: mean

src = Float32[1.0 0.5 2.0 1.5 0.8
0.7 1.2 0.9 1.1 1.3
1.4 0.6 1.8 0.4 1.0]
idx = [1, 2, 2, 3, 3]

f(op, x, idx) = sum(abs2, NNlib.scatter(op, x, idx))

for op in (+, -, *, /, max, min, mean)
dsrc = zeros(Float32, size(src))
try
Enzyme.autodiff(Reverse, f, Active, Const(op), Duplicated(copy(src), dsrc), Const(idx))
println(rpad(string(op), 5), " ok")
catch e
println(rpad(string(op), 5), " ", first(split(sprint(showerror, e), '\n')))
end
end
```

### Output

```
+ ok
- ok
* MethodError: no method matching reverse(::EnzymeCore.EnzymeRules.RevConfigWidth{1, false, true, (false, false, true, false, false), false, false}, ::Const{typeof(NNlib.scatter!)}, ::Type{DuplicatedNoNeed{Matrix{Float32}}}, ::Nothing, ::Const{typeof(*)}, ::Duplicated{Matrix{Float32}}, ::Duplicated{Matrix{Float32}}, ::Const{Vector{Int64}})
/ MethodError: no method matching reverse(... ::Const{typeof(/)} ...)
max MethodError: no method matching reverse(... ::Const{typeof(max)} ...)
min MethodError: no method matching reverse(... ::Const{typeof(min)} ...)
mean MethodError: no method matching reverse(... ::Const{typeof(mean)} ...)
```

The in-place `scatter!(op, dst, src, idx)` form fails for the same five ops.

Zygote returns correct gradients for all seven ops — I checked them against `FiniteDifferences.central_fdm(5, 1)` and they agree — so this only affects the Enzyme path.

### Where it bites

In GraphNeuralNetworks.jl this blocks Enzyme on every attention layer (they use `scatter` with `max` for the softmax) and on every mean aggregation, which is 8 layers in total.

### Versions

Julia 1.12.5, NNlib 0.9.44, Enzyme 0.13.199, Zygote 0.7.12.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the NNlib.scatter and scatter! entry points and run the provided Enzyme reproducer for all seven operations. Check the reverse-pass handling for the five failing operations, including the in-place form. Done means Enzyme no longer throws for those cases and the resulting gradients agree with the existing Zygote and finite-difference checks.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.