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

..

__init__.py

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

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

ad.py

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08:00

batching.py

Switch to new-style f-strings

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

mlir.py

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08:00

partial_eval.py

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08:00

pxla.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

xla.py

Don't call dtypes.result_type() unnecessarily on the type of an array during abstractification.

2022-11-29 19:27:10 -08:00