AutoEnzyme Jacobian preparation scalar-indexes Reactant traced arrays

[ChrisRackauckas-Claude]

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 #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

[gdalle]

I guess that's an issue to open on Enzyme itself, if Enzyme.jacobian doesn't like Reactant tracing, DI.jacobian won't either

[gdalle]

As discussed in #1067, an alternative fix would be to

• Use DI.AutoForwardFromPrimitive(AutoEnzyme())
• Lobby on the Enzyme side for Add chunk size selection strategies EnzymeAD/Enzyme.jl#2659, so that we have a way to specify that the whole input should be one chunk (I've been trying to get this merged for a year but I'm tiring out)

[ChrisRackauckas]

Kind of / maybe. I think here's a better question: is a good set of semantics here "DI.jacobian means Enzyme.jacobian inside Reactant compilation"? Because Reactant natively supports Enzyme in its compilation, and it really doesn't support any other AD, since any other AD mechanism is going to mutate in ways that won't trace well. So I think it's at least obvious that DI.jacobian on AutoEnzyme should translate to "just use the Reactant primitives it wants to use", but more generally, should it just be that "if you see DI.jacobian at all, ignore that AutoX choice and just make it use the Reactant primitives it wants to use"?

I can see arguments for both sides here. On one hand, if you decide to make everything Reactant-ified, then you're assuming it's taking over the compilation and just giving you whatever it thinks is fast (which is what Reactant tends to do), which means it "should" replace ForwardDiff with Enzyme stuff. On the other hand, if the user has stated "I want to use ForwardDiff", they could think it's breaking some contract if you replace it with ForwardDiff.

For practical reasons, I think it might be easiest to just give in and say "DI.jacobian means Enzyme.jacobian inside Reactant compilation" since the chance that ForwardDiff/Mooncake/etc. ever are supported inside of Reactant kernels is close to zero.

[gdalle]

The proper way to do this, eventually, will be support of AutoReactant(AutoEnzyme()) inside DI (#265), of which #918 is an early prototype (but I want to tackle it myself, not throw a bot at it, because the design is not obvious). I kind of planned to work on a wider repo revamp first, but if Reactant support is a priority for SciML uses I could make it a priority

[gdalle]

Oh ok re-reading your message I understand better what you mean. You want to transform any DI.jacobian into the Enzyme one during Reactant compilation. I find that a bit too opinionated, if only because finite difference Jacobians would also work within Reactant tracing

[ChrisRackauckas]

It is a priority, and I got most of it working already. I'm finding most of it is really just about making choices, it all "just works" you just have to disambiguate what you want the JIT compilation to actually do. SciML/LinearSolve.jl#1239 opts LinearSolve.jl itself into being primitives in the Reactant graph, because XLA has its own LU but that isn't good for reasons XYZ... so it just takes over which is the right thing for that library. And then higher level libraries like NonlinearSolve basically already "just work" with Enzyme because they already support general arrays, but SciML/NonlinearSolve.jl#1197 noted that the one piece that didn't work was DI, so I'm trying to find the right solution here.

[gdalle]

If Enzyme.jacobian becomes tracing-friendly (for instance by finally merging EnzymeAD/Enzyme.jl#2659 so that users can explicitly specify they want a single chunk), DI.jacobian tries to be as lightweight a wrapper as possible so it should work as well?

[ChrisRackauckas]

Oh ok re-reading your message I understand better what you mean. You want to transform any DI.jacobian into the Enzyme one during Reactant compilation. I find that a bit too opinionated, if only because finite difference Jacobians would also work within Reactant tracing

That's fine. The other approach is for the SciML libraries to say, okay we already have a system for choosing a default ADTypes method, and if we see Reactant, we choose Enzyme as the differentiation backend, in which case all we need as the contract is "DI.jacobian on AutoEnzyme means just do Enzyme.jacobian", which I think is reasonable. If the user chooses the wrong ADType, they get an error for trying to mix an AD that doesn't work with Reactant into Reactant. That would mean all DI needs to do is support Reactant for AutoEnzyme.

And yeah, it looks like the second half of this is... Enzyme.jacobian isn't supported in Reactant lol, and that seems to be upstream. I'm not 100% sure that fixing that will fix DI, but I'm going to take a stab at that today and talk with Billy about it.

[gdalle]

In any case, Billy basically forbade me to support AutoReactant(backend) for backend != AutoEnzyme so we got that covered ^^
A thornier question related to #265 is that DI's way of supporting Reactant will be "run Reactant.@compile during DI.prepare_jacobian", in other words, accelerate only the Jacobian (cause that's where DI's influence ends). Your use case is slightly different, whereby you want DI.jacobian to play nice inside a wider tracing procedure. I don't have a ready-made answer for that but I welcome suggestions

[ChrisRackauckas]

Another contract DI could do is turn its functions into Reactant primitives so they aren't traced, that would make something like AutoForwardDiff work even without ForwardDiff working, but then you'd lose the "Reactant-ification" of the internals there. I'm not sure if that's the right contract for this library.

In any case, Billy basically forbade me to support AutoReactant(backend) for backend != AutoEnzyme so we got that covered ^^

Okay, well I think that settles it then that then. I was trying to see how to get general NonlinearSolve code working, where lots of the time it's choosing AutoForwardDiff and such. But if the answer is "Billy thinks that case should error", then it's a matter of changing defaults handling too and erroring as expected.

Your use case is slightly different, whereby you want DI.jacobian to play nice inside a wider tracing procedure. I don't have a ready-made answer for that but I welcome suggestions

Yeah... but right now the issue seems to be in Enzyme itself

At least I see what semantics are wanted here now. There's some bug fixes to make that work though.