George Necula 2ccd785e16 Internal change ...

PiperOrigin-RevId: 514078303

2023-03-04 18:42:40 +00: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

Use math.prod instead of util.prod

2023-02-28 16:04:00 -08:00

jax2tf

Internal change

2023-03-04 18:42:40 +00:00

sparse

[sparse] modify impl rules to always use lowering path

2023-03-02 14:27:18 -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

Shorten alias chains for names exported in jax. namespace.

2023-03-01 09:19:44 -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

Internal change

2023-03-04 13:24:11 +00:00

x64_context.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00