Yash Katariya 1dd35f831c Add a multihost test for Array on non-continuous mesh ...

PiperOrigin-RevId: 464659865

2022-08-01 17:14:10 -07:00

..

__init__.py

[JAX] Move traceback_util.py into jax._src.

2020-11-04 09:02:59 -08:00

ad.py

fix ad_checkpoint.checkpoint caching issue

2022-07-29 19:59:28 -07:00

batching.py

[dynamic-shapes] start basic vmap compatibility

2022-07-09 10:03:40 -07:00

mlir.py

Enable debugging primitives in pjit on CPU/GPU

2022-07-29 20:10:27 -07:00

partial_eval.py

fix ad_checkpoint.checkpoint caching issue

2022-07-29 19:59:28 -07:00

pxla.py

Add a multihost test for Array on non-continuous mesh

2022-08-01 17:14:10 -07:00

xla.py

BUG: ensure that boolean scalars are never marked weak

2022-07-07 15:41:23 -07:00