Matthew Johnson cbcfe95e80 fix ad_checkpoint.checkpoint caching issue ...

Also add a config option to switch to the new checkpoint implementation
globally (default False for now), as the first step in replacing and then
deleting old remat.

2022-07-29 19:59:28 -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

Cache subjaxprs when lowering to MLIR functions.

2022-07-27 14:45:00 -07:00

partial_eval.py

fix ad_checkpoint.checkpoint caching issue

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

pxla.py

Merge pull request #11659 from mattjj:pmap-new-remat

2022-07-28 21:27:59 -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