Yash Katariya a4ca0dbc6c Make the signature of AbstractMesh to be AbstractMesh(axis_size: tuple[int, ...], axis_name: tuple[str, ...], *, axis_types) instead of AbstractMesh(shape_tuple: tuple[tuple[str, int], ...], *, axis_types) so that we are consistent across all Mesh APIs: Mesh, AbstractMesh and make_mesh ...

PiperOrigin-RevId: 736371111

2025-03-12 21:32:31 -07:00

..

__init__.py

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

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

_export.py

Make the signature of AbstractMesh to be AbstractMesh(axis_size: tuple[int, ...], axis_name: tuple[str, ...], *, axis_types) instead of AbstractMesh(shape_tuple: tuple[tuple[str, int], ...], *, axis_types) so that we are consistent across all Mesh APIs: Mesh, AbstractMesh and make_mesh

2025-03-12 21:32:31 -07:00

serialization_generated.py

support e2m1fn

2025-03-05 17:44:34 +00:00

serialization.fbs

support e2m1fn

2025-03-05 17:44:34 +00:00

serialization.py

support e2m1fn

2025-03-05 17:44:34 +00:00

shape_poly_decision.py

[shape_poly] Improve threefry with symbolic shapes

2025-01-07 09:10:04 +02:00

shape_poly.py

refactor: move jnp.einsum impl into its own submodule

2025-02-12 09:05:30 -08:00