Yash Katariya aa5e229027 Bump minimum jaxlib version to 0.4.4 which means xla_extension_version >= 127 ...

PiperOrigin-RevId: 512173011

2023-02-24 15:05:44 -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

Merge pull request #14615 from skye:restore_opt_barrier

2023-02-23 18:08:08 -08:00

sparse

[sparse] bcoo_dot_general_sampled: another special case

2023-02-24 10:50:54 -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

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

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

Merge pull request #14504 from skye:host_callback_pjrt_error

2023-02-15 17:11:01 -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

[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

Rename in_axis_resources and out_axis_resources with in_shardings and out_shardings. This is just a simple name replacement. It does not change any of the current pjit semantics and doesn't break any code.

2023-02-18 10:00:36 -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

[shard-map] add systematic tests for eager, jit, autodiff

2023-02-16 20:40:09 -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