Meta Issue for Reactant Compilation of Lux Models

[avik-pal]

WIP PRs

• Auto compile Lux models to reactant #665 -- Attempts to compile Lux models to XLA. While more generally scoped to allow models that cannot be compiled via Reactant, this is particularly hard because we need to give ownership of certain arrays to XLA (so high cost of data movement)
• feat: compile training loop automatically using reactant #969 -- compiles the training loop via single_train_step!. This should land fairly soon.

List of missing operations

It is mostly an incomplete list, but certain operations are known not to work

• conv
• batchnorm -- mean, std don't have bindings. fixed by feat: more coverage for common NN operations EnzymeAD/Reactant.jl#55

Shortcomings in Lux design that can be improved

• states are not updated in-place making it not ideal for both Enzyme and Reactant. This is honestly hard to fix at this point.

[wsmoses]

@avik-pal since Reactant has a lot more than just XLA, including custom linear algebra/fusion optimiations, MLIR AD, etc, can this be renamed Reactant compilation rather than xla?

[avik-pal]

Listing out models that would allow how to figure out how much of Lux we can compile atm

Target List of Models

Vision

• VGG
• ConvMixer -- https://github.com/LuxDL/Lux.jl/tree/main/examples/ConvMixer. Run this model with reactant and check perf
• VisionTransformer -- In Boltz.jl

Text

• Llama2

Other nice to haves

• Stable Diffusion

SciML

These are longer term goals because these need compiling NonlinearSolve/OrdinaryDiffEq/SciMLSensitivity which would be harder at this point

• Deep Equilibrium Networks -- DeepEquilibriumNetworks.jl
• NeuralODEs -- DiffEqFlux.jl
• PINNs -- Nested AD needed
• NeuralOperators -- we need fft support in Reactant (stablehlo has the ops)