Shaobo Hou 6e0c8029cc Fix type annotation for bcoo_update_layout. ...

PiperOrigin-RevId: 493567424

2022-12-07 04:25:25 -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

Merge pull request #13517 from froystig:rng-part-jax2tf

2022-12-06 08:46:28 -08:00

sparse

Fix type annotation for bcoo_update_layout.

2022-12-07 04:25:25 -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

Add typehints and point to the correct endpoint of Mesh and PartitionSpec in the args section.

2022-12-05 09:49:18 -08:00

ode.py

Merge pull request #11077 from DanPuzzuoli:ode_dt_max

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

pjit.py

Raise an error when a numpy input is passed with a non-trivial sharding. This can lead to weird behavior with pjit and XLA since host-local inputs are not allowed with pjit anymore.

2022-12-02 20:47:45 -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