jax
0417e1e5 - Fix `jax.lax.axis_index` in multi-host setting. (#1976)

Commit

5 years ago

Fix `jax.lax.axis_index` in multi-host setting. (#1976)

References

#1976 - Fix `jax.lax.axis_index` in multi-host setting.

Author

skye

skye

Parents

Loading