text: Title: keras.io examples using model.stateless_call aren't setting training=True and aren't managing non_trainable state
  Body: both the examples 
  * https://keras.io/guides/writing_a_custom_training_loop_in_jax/
  * https://keras.io/guides/distributed_training_with_jax/ 
  
  have a training loop that uses `stateless_call`, but neither sets `training=True` ( which means they are using the `training=False` behaviour )
  
  i guess it's not strictly relevant to the first example, which has no interesting non_trainables, but the second example has batch norm and dropout.
  
  additionally i don't see where this non_trainable state would be managed? 
  
  * re: batchnorm; each replica might have the same params, but is getting a different batch, so would collect different batch norm stats after the stateless_call. these aren't being aggregated?
  * re: dropout; the RNGKeys for dropout are in the non_trainable, but never change. shouldn't these be split after each forward pass?
  