mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 07:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src/export
History
George Necula 51dbcd4dad [export] Add backwards compatibility test for annotate_device_placement. 
This enables exporting functions that use memory kinds to place
data in different memories.

jax-fixit

PiperOrigin-RevId: 745008959
2025-04-08 00:10:23 -07:00
..
__init__.py
[export] Move the export implementation to jax._src.export.
2024-06-09 08:59:50 -07:00
_export.py
[export] Add backwards compatibility test for annotate_device_placement.
2025-04-08 00:10:23 -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
cleanup now that we depend on ml_dtypes>=0.5
2025-03-28 07:44:38 -07: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