rocm_jax/gda_serialization at 13e45c8953faa8e6fe13316ae61c06f16e94a322 - rocm_jax - Gitea For EOELAB

mirrors/rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-25 01:16:05 +00:00

History

Colin Gaffney 6ace6667dd Set coordinator address to allow it to later be used to initialize OCDBT coordinator server. Allow user to pass ts.Context when serializing or deserializing.

PiperOrigin-RevId: 520064049

2023-03-28 10:25:07 -07:00

..

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

README

Move serialization/de-serialization of GDA into jax.

2021-12-06 20:05:02 -08:00

serialization_test.py

Remove references to jax.config.jax_array, which is always True at head.

2023-03-15 17:09:11 -07:00

serialization.py

Set coordinator address to allow it to later be used to initialize OCDBT coordinator server. Allow user to pass ts.Context when serializing or deserializing.

2023-03-28 10:25:07 -07:00

README

# Serialization and De-serialization of GlobalDeviceArray via tensorstore

Warning: This directory is going to move in the near future. Please use at your
own risk.

To use this library, please install tensorstore and JAX.

```bash
pip install -U tensorstore
```