jax authors ed51c65576 Merge pull request #11405 from mattjj:djax-vmap ...

PiperOrigin-RevId: 459958155

2022-07-09 10:38:39 -07:00

..

__init__.py

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

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

ad.py

[dynamic shapes] revive iree

2022-07-06 15:01:16 -07:00

batching.py

[dynamic-shapes] start basic vmap compatibility

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

masking.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

mlir.py

Update minimum jaxlib version to 0.3.14.

2022-07-08 00:36:02 +00:00

partial_eval.py

[dynamic-shapes] revive basic bounded int machinery, add tests

2022-07-06 22:31:26 -07:00

pxla.py

Move _get_array_mapping from gda.py to pxla.py

2022-07-08 21:38:06 -07:00

sharded_jit.py

Enable Python callbacks on TFRT TPU backend

2022-07-06 20:52:50 -07:00

xla.py

BUG: ensure that boolean scalars are never marked weak

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