rocm_jax

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

History

George Necula a0812cd57e [better_errors] Make it explicit that debug_info is not None.

Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See https://github.com/jax-ml/jax/issues/26480 for more details.

PiperOrigin-RevId: 726770483

2025-02-13 22:07:04 -08:00

__init__.py

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

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

ad.py

[better_errors] Make it explicit that debug_info is not None.

2025-02-13 22:07:04 -08:00

batching.py

Rename sharding argument to out_sharding for lax.reshape, lax.broadcast_in_dim, lax.broadcast and lax.broadcasted_iota. .bind of these APIs still take sharding as a parameter though (but that's fine since it's internal and not public facing)

2025-02-12 13:59:23 -08:00

mlir.py

Added stream annotation support via @compute_on('gpu_stream:#')