Lux compatible layers

[ba2tripleO]

No description provided.

[ChrisRackauckas]

What's the status of this?

[ba2tripleO]

NeuralDSDE and NeuralSDE with Lux work now using EM or LambaEM with BacksolveAdjoint. But the AD vjp s are throwing some errors. It works only with numerical vjp right now. I think we'll need some changes to SciMLSensitivity to make vjp compatible?

[ChrisRackauckas]

How is it erroring? Use EulerHeun so it's in Stratanovich form: that's simpler for training.

[ChrisRackauckas]

@avik-pal is that definition correct? https://github.com/SciML/DiffEqFlux.jl/actions/runs/3855851152/jobs/6571364733#step:6:777

[avik-pal]

I had never stored parameters in the layers so never tested @functor. The fallback definition there is to allow recursively traverse container layers.

For Flux compat, I think we should directly define the Optimisers.trainable though that will not allow Flux.gpu to work as expected.

[ChrisRackauckas]

But it seems it shouldn't dispatch on the second argument if it's going to work with @functor?

[avik-pal]

LuxDL/Lux.jl#229 should fix it. I will tag a release once the CI passes

[ChrisRackauckas]

@Abhishek-1Bhatt https://github.com/SciML/DiffEqFlux.jl/actions/runs/3855851152/jobs/6581195485 this failure looks real, but looks like it just needs a functor overload.

[ba2tripleO]

The error message, seems to be referring to model as in NeuralDELayer <: Lux.AbstractExplicitContainerLayer{(:model,)} as HamiltonianNeuralDE is also subtyped to NeuralDELayer here , maybe doing something like this could help?

[ChrisRackauckas]

yes the naming seems wrong, so it either needs to rename to just model or have another abstract type. I would recommend just naming it model like the others, I think that might make the structure simpler.