On: Training Deep Networks with Data Parallelism in Ja...

[mccoyb]

I would start here: https://jax.readthedocs.io/en/latest/notebooks/Writing_custo...

Note that JAX enforces certain limitations, so you should be careful when considering JAX to be "the perfect language" - in general I don't think this is true. It's quite good at what it's designed for.

[credit_guy → mccoyb]

Thank you