Use `Lux.Experimental.StatefulLuxLayer` internally to avoid `Box`ing

[avik-pal]

See https://lux.csail.mit.edu/dev/tutorials/intermediate/1_NeuralODE#type-stability for the issue with the current implementation style.