Add Reactant support

[wsmoses]

Using reactant to compile the inner update loop will likely result in a significant performance boost (removing intermediate allocations/computations), in addition to auto upgrading all code to GPU/TPU where possible

cc @avik-pal

[avik-pal]

Some starter code https://github.com/LuxDL/Lux.jl/blob/main/ext/LuxReactantExt/training.jl

[avik-pal]

We also added reactant patches to make ComponentArrays work nicely, so it should also help the neural network cases if integrated with Optimization.

For API, I would recommend reusing AutoEnzyme with get_device_type from MLDataDevices to use Reactant if the device is ReactantDevice.

[gdalle]

Could this be the opportunity to revisit Reactant support in DI? JuliaDiff/DifferentiationInterface.jl#325 was closed by Billy but I do think it would make sense to provide a reactant tracing mechanism as part of DI preparation, and I'd love to give it a second shot

[avik-pal]

For optimization.jl we should have a separate reactant ext even if DI has support for it. We will be able to optimize the whole inner loop which might get eliminate more intermediates. See for example how Lux has a return gradients option specifically for reactant

[wsmoses]

from the issue (you linked to a PR which you closed, this is the issue of relevance: JuliaDiff/DifferentiationInterface.jl#265)

"""
it might be, but it's also not really the right abstraction (e.g. you ought be able to @compile outside of any DI call as it has nothing to do with AD).

Of course it makes it easier to use via DI and enables the use of Enzyme [per other discussions], but also it's not really critical so I closed.
"""

And I think this remains true.

Having DI be able to do reactant things is useful for people wanting to try stuff out, especially if they want to swap out the derivative compiutation.

That also said if your code is more than just a derivative (e.g. x += gradient), you'd want to compile the entire code, not just the derivative which wouldn't be possible to do with a compile within DI.

In this case, like Avik mentioned optimizing the whole inner loop can likely remove additional intermediates (at minimum because an in place update means we don't need to allocate the intermediate buffers).

So, @gdalle definitely feel free to go for it, it will unquestionably be useful (especially for people wanting to try things out/test EnzymeMLIR), but also as a forewarning most end users probably won't want to use it through DI since they'd want to compile the entire computation, not just derivative. That said if plumbing is annoying, they could use it anyways.

[gdalle]

Ok I think I understand, and I agree that ideally people will want to compile the whole loop. But that would require them actually learning how to use Reactant, creating traced array, compiling, etc. In a code where the derivative is the most costly part of the hot loop, having an AutoReactant(AutoEnzyme()) which automatically compiles away the derivative during preparation could possibly help a lot of beginners. So I still think I'll give it a shot now that Reactant has evolved a bit more! Excited to try it out again

[wsmoses]

In fairness Reactant explicitly is trying to be easy to use (eg just convert your array types, and call compile and or jit), but it’ll definitely be useful for all the code where AD is critical for folks to start using reactant.

not a discussion for this issue, but I think there is a broader design choice to made here for DI. In particular I see Reactant much more like CUDA.jl in the sense that there’s a reactant array you use or a cuda array you use (which is why lux uses this from the data type). So I’m not sure an ADType makes sense, but you could have some data type wrapper and DI takes in an object that can encapsulate both?

[Vaibhavdixit02]

This will have to be done both as an extension here and modifications to the solver libraries iiuc. In which case I think it will be usable with a small subset of the solver libraries in which we have access to the inner loop. I will start experimenting.