Allow embedded neural network evaluation to have side effects

[jrmaddison]

Required e.g. for batch normalization

[jrmaddison]

The simplest way to do this might be to register Keras Layers as pytree nodes, repurposing Keras serialization for the flatten and unflatten, and then swap the vmap and fori_loop used to evaluate the Dynamics layer so that the embedded network operates on the whole batch. It might alternatively be possible to have Dynamics inherit from keras.layers.RNN.

[jrmaddison]

Keras cannot deserialize when Traced variables appear, which I think would be needed for the pytree node approach.