Enzyme: scatter! reverse rule only supports + and -, other ops throw MethodError
- 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