jax authors c15f30f22e Instrument a new metric to measure the savings of async checkpoint in JAX. ...

Create the new metric '/jax/checkpoint/write/async/thread_duration_sec' to measure the savings from the async thread creation time.

PiperOrigin-RevId: 529227213

2023-05-03 16:36:54 -07:00

..

__init__.py

Rename gda_serialization to array_serialization but keep gda_serialization around until it is included in a jax release so that OSS projects can be moved to array_serialization

2023-03-31 18:07:51 -07:00

serialization_test.py

Change to simpler import for jax.config

2023-04-21 11:51:22 -07:00

serialization.py

Instrument a new metric to measure the savings of async checkpoint in JAX.

2023-05-03 16:36:54 -07:00