Yash Katariya 52a7701dda Replace usage of {in|out}_axis_resources with {in|out}_shardings ...

PiperOrigin-RevId: 513040164

2023-02-28 14:29:09 -08:00

..

compilation_cache

Bump minimum jaxlib version to 0.4.4 which means xla_extension_version >= 127

2023-02-24 15:05:44 -08:00

gda_serialization

Rename jax.sharding.OpShardingSharding to jax.sharding.GSPMDSharding. jax.sharding.OpShardingSharding will be removed in 3 months from Feb 17, 2023.

2023-02-17 17:11:06 -08:00

jax2tf

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -08:00

sparse

Refactor bcoo_dot_general GPU lowering

2023-02-28 13:36:52 -08:00

__init__.py

Change jax.ad, jax.xla, jax.pxla to point to the shims instead of the internal modules.

2023-02-09 13:31:40 -08:00

checkify.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

custom_partitioning.py

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -08:00

global_device_array.py

Move contents of jax.experimental.global_device_array to jax._src.global_device_array.

2023-02-16 15:37:10 -08:00

host_callback.py

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -08:00

jet.py

custom_jvp symbolic zeros support

2023-02-24 07:33:49 -08:00

maps.py

[JAX] Deprecate jax.experimental.maps.Mesh.

2023-02-15 09:15:50 -08:00

mesh_utils.py

Replace uses of deprecated JAX sharding APIs with their new names in jax.sharding.

2023-02-03 14:28:45 -08:00

multihost_utils.py

Replace usage of {in|out}_axis_resources with {in|out}_shardings

2023-02-28 14:29:09 -08:00

ode.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

pjit.py

Deprecate uses of jax.experimental.pjit.NamedSharding and jax.experimental.pjit.PartitionSpec.

2023-02-15 08:14:26 -08:00

rnn.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

serialize_executable.py

Add support for a list of allow_spmd_sharding_propagation_to_output. This gives us more flexibility to tell SPMD which shardings to override.

2023-02-03 17:59:10 -08:00

shard_map.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -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