rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00

History

Peter Hawkins 262a4f482c Deprecate support for custom lowering rules that return tuple-wrapped ir.Values.

https://github.com/google/jax/pull/22211 forbade custom lowering rules from returning singleton tuples of ir.Value, but this appears to break downstream users, notably Transformer Engine. Instead, allow lowering rules to return singleton tuples and unwrap them if needed, but warn if this behavior is seen.

PiperOrigin-RevId: 650345051

2024-07-08 12:54:44 -07:00

__init__.py

Reapply: move jax.interpreters.ad to jax._src.interpreters.ad

2023-02-02 09:29:05 -08:00

ad.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

batching.py

Update jax/_src/api.py

2024-07-02 16:10:36 +01:00

mlir.py

Deprecate support for custom lowering rules that return tuple-wrapped ir.Values.