jax authors b467feb250 [JAX] Add RunTimeError to host_local_array_to_global_array ...

PiperOrigin-RevId: 485657586

2022-11-02 12:06:26 -07:00

..

compilation_cache

Increase the minimum jaxlib version to 0.3.22.

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

gda_serialization

jax.Array is a unified type that will subsume JAX's DeviceArray, ShardedDeviceArray and GlobalDeviceArray.

2022-10-14 14:09:01 -07:00

jax2tf

Create a table of contents for jax2tf documentation.

2022-11-01 00:35:49 -07:00

sparse

[sparse] fix a typo in n_sparse calculation.

2022-11-01 11:46:10 -07:00

__init__.py

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

2022-09-22 12:27:19 -07: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

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

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

custom_partitioning.py

Increase the minimum jaxlib version to 0.3.22.

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

global_device_array.py

Fix docstring in GDA

2022-11-01 08:36:06 -07:00

host_callback.py

Fix typo.

2022-10-14 22:12:08 -07:00

jet.py

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

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

maps.py

Don't assume that vmap doesn't introduce constants

2022-11-02 08:27:10 -07:00

mesh_utils.py

[mesh_utils] Support single-core 2D meshes

2022-10-26 11:32:50 -07:00

multihost_utils.py

[JAX] Add RunTimeError to host_local_array_to_global_array

2022-11-02 12:06:26 -07:00

ode.py

Merge pull request #11077 from DanPuzzuoli:ode_dt_max

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

pjit.py

Add the full shape to the error message too (in pjit_check_aval_sharding) to give users full information about what is going on.

2022-10-28 09:14:11 -07:00

x64_context.py

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

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