Hyeontaek Lim 218305964f Fix GDA error message formatting ...

PiperOrigin-RevId: 486724647

2022-11-07 11:55:28 -08:00

..

compilation_cache

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

2022-11-02 17:28:38 -07: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

[jax2tf] Fixed jax2tf Limitations

2022-11-07 11:26:39 +02:00

sparse

[sparse] simplify BCOO indexing implementation

2022-11-02 10:13:35 -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 GDA error message formatting

2022-11-07 11:55:28 -08:00

host_callback.py

Add a number of missing function cross-references in the docs.

2022-11-07 12:00:26 -05:00

jet.py

BUG: fix jet rule for dynamic_slice

2022-11-02 21:37:41 -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

If the input to host_local_array_to_global_array is not fully addressable (i.e. not host local), return it as is.

2022-11-05 20:16:14 -07:00

ode.py

Merge pull request #11077 from DanPuzzuoli:ode_dt_max

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

pjit.py

Remove static_argnums from AOT invocation.

2022-11-07 10:21:57 -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