Kostiantyn Liepieshov c730bbda74 fix bug in export_module when no mesh axes are empty for shardy. ...

If mesh axes are empty, we are setting mesh as None, resulting in an error in
this test.

This fix provides an empty mesh, when no mesh axes in dumped module are empty.

PiperOrigin-RevId: 746058506

2025-04-10 09:21:58 -07:00

..

__init__.py

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

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

_export.py

fix bug in export_module when no mesh axes are empty for shardy.

2025-04-10 09:21:58 -07:00

serialization_generated.py

[export] Add support for serializing functions with PRNG keys as inputs/outputs

2025-04-07 11:53:20 +02:00

serialization.fbs

[export] Add support for serializing functions with PRNG keys as inputs/outputs

2025-04-07 11:53:20 +02:00

serialization.py

[export] Add support for serializing functions with PRNG keys as inputs/outputs

2025-04-07 11:53:20 +02:00

shape_poly_decision.py

[shape_poly] Improve threefry with symbolic shapes

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

shape_poly.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00