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.

Upcoming Work

• Compiling Lux models and offload part of the computation to XLA and adding ChainRules and EnzymeRules support to call the generated VJP rules

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.

Core Layers Missing Adjoints

• Embedding -- stablehlo.dynamic_gather
• Pooling -- stablehlo.reduce_window
• Conv -- stablehlo.convolution / stablehlo.dynamic_conv

[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?