AutoEnzyme Jacobian preparation scalar-indexes Reactant traced arrays

Problem

prepare_jacobian with AutoEnzyme() fails when called inside Reactant.@jit because Enzyme's chunked shadow preparation scalar-indexes a TracedRArray.

This is a focused instance of the broader Reactant support discussion in https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/265.

Reproducer

This environment contains only DifferentiationInterface, Enzyme, and Reactant; it imports no SciML package.

using DifferentiationInterface: AutoEnzyme, jacobian, prepare_jacobian
using Enzyme: Enzyme
using Reactant: Reactant, @jit

f(x) = x .^ 2

function di_jacobian(x)
    backend = AutoEnzyme()
    prep = prepare_jacobian(f, backend, x)
    return jacobian(f, prep, backend, x)
end

x = Reactant.to_rarray(Float32[1, 2])
@jit di_jacobian(x)

The call fails during tracing:

ERROR: LoadError: Scalar indexing is disallowed.
Invocation of setindex!(::TracedRArray, v, ::Union{Int, TracedRNumber{Int}}) resulted in scalar indexing of a GPU array.
...
 [16] #onehot##2
    @ Enzyme/src/sugar.jl:153
 [19] #chunkedonehot##0
    @ Enzyme/src/sugar.jl:424
 [25] create_shadows
    @ Enzyme/src/sugar.jl:436
 [28] prepare_jacobian_nokwarg
    @ DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl:241

The equivalent public Enzyme primitive works inside the same Reactant kernel:

function enzyme_jacobian(x)
    derivatives = Enzyme.autodiff(
        Enzyme.Forward,
        f,
        Enzyme.BatchDuplicated(x, Enzyme.onehot(x)),
    )
    return stack(values(only(derivatives)))
end

@jit enzyme_jacobian(x)
# Reactant.ConcretePJRTArray(Float32[2 0; 0 4])

Delegating the traced path to this primitive would also avoid DI's scalar-indexing basis construction for in-place functions. The higher-level Enzyme.jacobian convenience function cannot currently be used directly: its forward implementation calls tupstack, which has no traced-array method, while its reverse implementation fails during its augmented forward pass.

Expected behavior

prepare_jacobian and jacobian with the default AutoEnzyme() should work inside Reactant.@jit and return a Reactant array.

Versions

Julia Version 1.12.6
Commit 15346901f00 (2026-04-09 19:20 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 128 × AMD EPYC 7502 32-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-18.1.7 (ORCJIT, znver2)
  GC: Built with stock GC
Threads: 1 default, 1 interactive, 1 GC (on 128 virtual cores)

DifferentiationInterface v0.7.21 (main at f0fb136d)
Enzyme v0.13.199
Reactant v0.2.283

Run the supplied Julia reproducer with Reactant.@jit, then inspect DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl around line 241 and Enzyme/src/sugar.jl around the reported shadow-preparation calls. Compare the traced path with the working Enzyme.autodiff primitive and verify that prepare_jacobian and jacobian return a Reactant array without scalar indexing.