Complete Lux+Reactant migration: inference, callbacks, and multi-loss support

[AdityaPandeyCN]

This PR adds _forward using Lux.apply with separated params/state for inference, _sync_params_to_model! to copy TrainState to CPU for callbacks, _get_lux_loss to dispatch the correct loss per task, one-hot encoding for classification targets, and fit_iter! for MLJ support.

[AdityaPandeyCN]

Please take a look when free @jeremiedb.

[jeremiedb]

What is the need for all_batches?
Why not just keep dtrain_loader? Collecting the iterator doesn't seem a sound practice, it would notably cause problems if the dataloader was a lazy iterator accessing larger than memory data.
Also, I wasn't aware of of any issue when using Flux, nor in the initial Lux adaption under neuro-v3.

As per the one-hot, I don't think this should be needed either, as I'd expect the crossentropy loss to work fine with an integer encoding of the multi-class target. Or is it something specific that Lux doesn't support while Flux does?

[jeremiedb]

The idea for stack tree was to facilitate the stacking of stack_size NeuroTrees in a residual network fashion.
That being said, maybe that model contruction in Lux makes it easy to do the equivalent by composing NeuroTrees in the same way as a vanilla residual architecture would do with Dense layers?

[jeremiedb]

Thanks! The initial benchmark look good.
I appreciate the cleanup regarding code dupliocation in inference, it was long due!

I left a couple of comments, which I think should be striaghforward to adress.
There's the collect of batches which appeared at first sight as not being a change needed, or that may be need a redesign within the DataLoader logic.