Yash Katariya 934bc4e1b3 Move PartitionSpec and Mesh out of experimental and into the sharding namespace. The new API endpoint is jax.sharding.PartitionSpec and jax.sharding.Mesh. ...

PiperOrigin-RevId: 492358238

2022-12-01 19:28:32 -08:00

..

compilation_cache

Add c++ pickling for all the jax Sharding types.

2022-11-02 17:28:38 -07:00

gda_serialization

Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL).

2022-11-14 14:44:00 -08:00

jax2tf

Switch to new-style f-strings

2022-12-01 09:14:16 -08:00

sparse

Switch to new-style f-strings

2022-12-01 09:14:16 -08:00

__init__.py

Move PartitionSpec and Mesh out of experimental and into the sharding namespace. The new API endpoint is jax.sharding.PartitionSpec and jax.sharding.Mesh.

2022-12-01 19:28:32 -08:00

callback.py

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

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

checkify.py

Rewrite Checkify to support tracking different error types.

2022-11-25 15:31:54 +00:00

custom_partitioning.py

Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL).

2022-11-14 14:44:00 -08:00

global_device_array.py

Remove local imports of array.py. The remaining local imports are in pxla.py but I will chip away at them when we delete SDA and move some more APIs out of experimental.

2022-11-30 15:26:03 -08:00

host_callback.py

[x64] better type safety for host_callback

2022-12-01 11:47:07 -08:00

jet.py

BUG: fix jet rule for dynamic_slice

2022-11-02 21:37:41 -07:00

maps.py

Move PartitionSpec and Mesh out of experimental and into the sharding namespace. The new API endpoint is jax.sharding.PartitionSpec and jax.sharding.Mesh.

2022-12-01 19:28:32 -08:00

mesh_utils.py

Generalize TPU mesh computations.

2022-11-21 03:22:24 -08:00

multihost_utils.py

Optimize host_local_array_to_global_array by caching the local to global conversion and flattening of axis resources. Also take a fast path for device_put which does not do abstractify and only canonicalize_dtype on the entire array once (instead of doing it for every shard).

2022-11-20 20:53:02 -08:00

ode.py

Merge pull request #11077 from DanPuzzuoli:ode_dt_max

2022-10-28 16:05:25 -07:00

pjit.py

Fix vmap(jvp(pjit(f))) when pjit doesn't have any axis_resources

2022-12-01 10:42:26 -08:00

rnn.py

Add CUDNN custom call for LSTM. Exposed as jax.experimental.rnn module.

2022-11-28 14:31:48 -08:00

serialize_executable.py

Refactor create_cpp_call to be a method on MeshExecutable

2022-11-17 18:05:10 -08:00

x64_context.py

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

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