[martin-gorner]

Support for distributed JAX training demoed here: bit.ly/keras-on-jax-demo You have to write a custom training loop for now, but it works.

[dharmeshkakadia]

Thanks!