Zafarali Ahmed 6e00ba8bad Enable more mesh shape assignment ...

We now sort the mesh dims by size first. Smaller dims have fewer choices so
they should be assigned first.

PiperOrigin-RevId: 520942700

2023-03-31 09:36:16 -07:00

..

compilation_cache

expose compiler_options on compile()

2023-03-30 17:14:26 -07:00

gda_serialization

Set coordinator address to allow it to later be used to initialize OCDBT coordinator server. Allow user to pass ts.Context when serializing or deserializing.

2023-03-28 10:25:07 -07:00

jax2tf

[shape_poly] Extend the handling of jnp.arange with shape polymorphism.

2023-03-31 14:41:26 +02:00

sparse

Merge pull request #15243 from jakevdp:coo-sort-warning

2023-03-27 16:50:24 -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

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

host_callback.py

Bump minimum jaxlib version from 0.4.6 to 0.4.7.

2023-03-28 13:43:01 -07:00

jet.py

Move jax.interpreters.partial_eval to jax._src.interpreters.partial_eval.

2023-03-27 13:30:47 -07:00

maps.py

Remove _PositionalSemantics class since it is not used anymore because jax.Array always has GLOBAL semantics

2023-03-17 13:30:04 -07:00

mesh_utils.py

Enable more mesh shape assignment

2023-03-31 09:36:16 -07:00

multihost_utils.py

Add deprecation warning for FROM_GDA usage since that argument is not required anymore.

2023-03-27 11:33:11 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

pjit.py

Delete unused functions from dispatch.py and pjit.py

2023-03-30 13:38:44 -07:00

rnn.py

Copy seq_lengths before creating descriptor

2023-03-27 10:59:44 -07:00

serialize_executable.py

Redefine compile_and_serialize as serialize(lowered.compile()).

2023-03-22 17:23:19 -07:00

shard_map.py

[shard-map] handle closed-over vmap tracers

2023-03-30 16:43:40 -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