rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-24 20:56:06 +00:00

History

George Necula 292a00b35a [export] Cleanup in the export module.

With jax.experimental.export gone we can now do some cleanup in the export module.

In particular we remove the `export.args_spec` API, and the `lowering_platforms` arg for `export.export`. These were deprecated in June 2024.

PiperOrigin-RevId: 692398132

2024-11-01 22:56:44 -07:00

__init__.py

[export] Move the export implementation to jax._src.export.

2024-06-09 08:59:50 -07:00

_export.py

[export] Cleanup in the export module.

2024-11-01 22:56:44 -07:00

serialization_generated.py

[export] Add support for serialization for some custom PyTree nodes

2024-10-21 11:38:13 +02:00

serialization.fbs

[export] Cleanup in the export module.