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

google / uncertainty-baselines

High-quality implementations of standard and SOTA methods on a variety of tasks.

Apache License 2.0

1.46k stars 203 forks source link

Closed copybara-service[bot] closed 1 year ago

copybara-service[bot] commented 1 year ago

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

This change updates:

{jax.experimental.maps.Mesh, jax.interpreters.pxla.Mesh} to jax.sharding.Mesh
{jax.experimental.PartitionSpec, jax.experimental.pjit.PartitionSpec, jax.interpreters.pxla.PartitionSpec, jax.pxla.PartitionSpec} to jax.sharding.PartitionSpec
jax.experimental.maps.NamedSharding to jax.sharding.NamedSharding.