Yash Katariya d02f28199b Clean up pjit after jax.Array ...

* Remove {in|out}_positional_semantics from pjit_p.bind
* Remove `in_is_global` from lower_sharding_computation
* Remove local_to_global and global_to_local
* Clean up some arguments of sharded_lowering since they are not needed

PiperOrigin-RevId: 517469390

2023-03-17 11:53:00 -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

Remove references to jax.config.jax_array, which is always True at head.

2023-03-15 17:09:11 -07:00

jax2tf

Clean up pjit after jax.Array

2023-03-17 11:53:00 -07:00

sparse

Clean up pjit after jax.Array

2023-03-17 11:53:00 -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

Improve pytype inference for Sharding type.

2023-03-14 14:20:17 -07:00

host_callback.py

Shorten alias chains for names exported in jax. namespace.

2023-03-01 09:19:44 -08:00

jet.py

Clean up pjit after jax.Array

2023-03-17 11:53:00 -07: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

[Jax cleanup]

2023-03-16 15:47:28 -07:00

ode.py

Remove leading underscores in jax._src.numpy.util

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

pjit.py

Clean up pjit after jax.Array

2023-03-17 11:53:00 -07:00

rnn.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

serialize_executable.py

Remove references to jax.config.jax_array, which is always True at head.

2023-03-15 17:09:11 -07: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