Jake VanderPlas 7ab7b214ac refactor: move jnp.einsum impl into its own submodule

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

..

__init__.py

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

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

_export.py

Remove old backward compatibility mode for old PRGN custom call on GPU

2025-02-04 07:34:52 -08:00

serialization_generated.py

Add e8m0fnu support by conditional dtype.

2025-01-22 21:57:43 +00:00

serialization.fbs

Add e8m0fnu support by conditional dtype.

2025-01-22 21:57:43 +00:00

serialization.py

Add e8m0fnu support by conditional dtype.

2025-01-22 21:57:43 +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