Yash Katariya ca1f58e37b Add a new jax.spmd_mode config for preventing unintentional hangs and incorrect results when users pass jax.Arrays that span across multiple processes (i.e. not fully addressable) to jit or jnp operations (that are jitted by default). ...

Implicitly jitted functions will **always** require a `jax.spmd_mode` context manager for operating on non-fully addressable jax.Array.

Explicitly jitted functions will require the `jax.spmd_mode` config to begin with as we roll out jax.Array since its a new behavior for `jit` (previously jit only worked on single device arrays).
* Overtime (via docs) and as users become more familiar with the new parallelism APIs, we can relax this restriction and allow explicit `jit` to work without needing the config. This can happen when we merge the frontend of `jit` and `pjit`.

PiperOrigin-RevId: 485075693

2022-10-31 09:51:42 -07:00

..

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

ad.py

[dynamic-shapes] fix nested vmap callable annotation logic

2022-10-19 17:01:53 -07:00

batching.py

make leak checker errors explain why objects are alive

2022-10-28 14:12:17 -07:00

mlir.py

Increase the minimum jaxlib version to 0.3.22.

2022-10-27 10:24:11 -04:00

partial_eval.py

make leak checker errors explain why objects are alive

2022-10-28 14:12:17 -07:00

pxla.py

Add a new jax.spmd_mode config for preventing unintentional hangs and incorrect results when users pass jax.Arrays that span across multiple processes (i.e. not fully addressable) to jit or jnp operations (that are jitted by default).

2022-10-31 09:51:42 -07:00

xla.py

Include column information in Python locations under Python 3.11.

2022-10-14 19:14:35 +00:00