mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 15:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/interpreters
History
Peter Hawkins f62ab8fef0 Use MLIR bytecode when passing IR to backends. 
MLIR bytecode is more compact to represent and should be faster to generate and parse.

PiperOrigin-RevId: 472701621
2022-09-07 05:41:06 -07:00
..
__init__.py
[JAX] Move traceback_util.py into jax._src.
2020-11-04 09:02:59 -08:00
ad.py
Merge pull request #11900 from sharadmv:for-loop
2022-08-18 20:14:17 -07:00
batching.py
Add spmd_axis_name to vmap to allow constraining mapped PartitionSpecs.
2022-08-08 19:41:42 -07:00
mlir.py
Use MLIR bytecode when passing IR to backends.
2022-09-07 05:41:06 -07:00
partial_eval.py
Add partial_eval_custom rule for for_loop
2022-09-06 11:00:26 -07:00
pxla.py
Add support to handle arbitrary shardings to KeyArray. Resolve all the TODOs that were created before.
2022-08-31 22:54:06 -07:00
xla.py
BUG: ensure that boolean scalars are never marked weak
2022-07-07 15:41:23 -07:00