Roy Frostig 219723c738 migrate internal dependencies from jax.interpreters.ad to jax._src.interpreters.ad ...

... in preparation for paring down `jax.interpreters.ad`'s exported symbols.

Includes some import fixups along the way.

PiperOrigin-RevId: 507684262

2023-02-06 22:52:36 -08:00

..

compilation_cache

Correctly hash auto_spmd fields in compilation cache key.

2023-02-06 17:15:23 -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

migrate internal dependencies from jax.interpreters.ad to jax._src.interpreters.ad

2023-02-06 22:52:36 -08:00

sparse

migrate internal dependencies from jax.interpreters.ad to jax._src.interpreters.ad

2023-02-06 22:52:36 -08:00

__init__.py

Add IO callback

2023-01-17 13:55:05 -08:00

checkify.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

custom_partitioning.py

Add static_argnums to custom_partitioning.

2023-02-03 11:41:17 -08:00

global_device_array.py

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

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

host_callback.py

Bump minimum jaxlib_version to 0.4.2 i.e xla_extension_version == 119 and mlir_api_version == 43

2023-02-06 10:37:33 -08:00

jet.py

Add pjit support in jax.experimental.jet

2023-01-23 15:51:47 -08:00

maps.py

Silence some pytype errors.

2023-01-30 15:33:20 -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

If there is only 1 process in process_allgather then just pull it to host without going via pjit.

2023-02-05 14:01:21 -08:00

ode.py

fix transcription error for initial_step_size (thanks @AlbertMitjans)

2023-01-14 21:10:44 -08:00

pjit.py

Move pjit.py to jax/_src in preparation for merging the jit and pjit frontend APIs

2022-12-16 13:07:15 -08:00

rnn.py

CI: fix mypy jaxlib version

2023-01-24 06:57:23 -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

Merge pull request #14323 from mattjj:shmap-add-trivial-rules

2023-02-06 18:23:01 -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