jax authors 31aee3ae70 Merge pull request #14954 from mattjj:ravel-pytree-jitfriendly ...

PiperOrigin-RevId: 516332491

2023-03-13 15:01:58 -07:00

..

compilation_cache

Fix tests that were adding __annotations__ for some reason to CompileOptions.

2023-03-08 13:55:03 -08:00

gda_serialization

Split Sharding subclasses out of _src/sharding.py into _src/sharding_impls.py

2023-03-13 08:50:18 -07:00

jax2tf

[jax2tf] Refactor the native lowering

2023-03-13 16:39:45 +02:00

sparse

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07: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

Split Sharding subclasses out of _src/sharding.py into _src/sharding_impls.py

2023-03-13 08:50:18 -07: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

Split Mesh and ResourceEnv into a new module jax._src.mesh.

2023-03-10 10:08:21 -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

Split Sharding subclasses out of _src/sharding.py into _src/sharding_impls.py

2023-03-13 08:50:18 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07: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

use Partial to make ravel_pytree unflatteners jit-friendly

2023-03-13 11:06:56 -07:00

x64_context.py

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

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