Missing adjoint of stablehlo.gather

[yolhan83]

Hello, I made an mwe for this,

using Enzyme,Reactant
Reactant.set_default_backend("cpu")
function foo(x,ids)
	return sum(x[ids])
end
function get_grad(x,ids)
	dx = Enzyme.make_zero(x)
	Enzyme.autodiff(Enzyme.Reverse, foo, Enzyme.Duplicated(x,dx), Enzyme.Const(ids))
	return dx	
end
dev= Reactant.to_rarray
x = rand(1000) |> dev
ids = rand(1:1000,100)
L = @compile foo(x,ids) # works
@info L(x,ids) # works
g = @compile get_grad(x,ids) # fail
@info g(x,ids)

I get the following error,

error: could not compute the adjoint for this operation %2 = "stablehlo.gather"(%arg0, %1) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1>}> : (tensor<1000xf64>, tensor<100x1xi64>) -> tensor<100x1xf64>
ERROR: "failed to run pass manager on module"
Stacktrace:
 [1] run!
   @ ~/.julia/packages/Reactant/nKvpM/src/mlir/IR/Pass.jl:79 [inlined]
 [2] run_pass_pipeline!(mod::Reactant.MLIR.IR.Module, pass_pipeline::String; enable_verifier::Bool)
   @ Reactant.Compiler ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:343
 [3] run_pass_pipeline!
   @ ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:338 [inlined]
 [4] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{ConcreteRArray{Float64, 1}, Vector{Int64}}; optimize::Bool, no_nan::Bool, backend::String)
   @ Reactant.Compiler ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:487
 [5] compile_mlir!
   @ ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:439 [inlined]
 [6] compile_xla(f::Function, args::Tuple{ConcreteRArray{Float64, 1}, Vector{Int64}}; client::Nothing, optimize::Bool, no_nan::Bool, device::Nothing)
   @ Reactant.Compiler ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:1007
 [7] compile_xla
   @ ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:986 [inlined]
 [8] compile(f::Function, args::Tuple{ConcreteRArray{Float64, 1}, Vector{Int64}}; sync::Bool, kwargs::@Kwargs{client::Nothing, no_nan::Bool, device::Nothing, optimize::Bool})     
   @ Reactant.Compiler ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:1055
 [9] top-level scope
   @ ~/.julia/packages/Reactant/nKvpM/src/Compiler.jl:709

on Enzyme v0.13.28 and Reactant v0.2.24

[mofeing]

The issue is that we don't yet have a diff rule for stablehlo.gather as stated in EnzymeAD/Enzyme-JAX#88