Peter Hawkins 00d45feee6 Deprecate uses of jax.experimental.pjit.NamedSharding and jax.experimental.pjit.PartitionSpec. ...

Use the aliases under jax.sharding instead.

PiperOrigin-RevId: 509837529

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

..

compilation_cache

Remove silent data corruption runtime flags from persistent cache key.

2023-02-08 12:31:27 -08:00

gda_serialization

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

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

jax2tf

Merge pull request #14469 from gnecula:poly_percentile

2023-02-15 07:33:26 -08:00

sparse

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

2023-02-14 23:01:11 -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

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

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

global_device_array.py

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

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

host_callback.py

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

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

jet.py

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

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

maps.py

Minimize exported names from jax.experimental.maps.

2023-02-13 12:57:54 -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

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

2023-02-14 23:01:11 -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

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

2023-02-14 23:01:11 -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

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

2023-02-14 23:01:11 -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