Yash Katariya 5b808c93d0 Make Shard a normal class making index and replica_id properties on that class. Raise an error if an indices cannot be calculated from a sharding. ...

PiperOrigin-RevId: 454899275

2022-06-14 10:37:19 -07:00

..

compilation_cache

Using etils(gfile) to support gcs buckets and file system for persistent compilation caching

2022-06-10 00:17:13 +00:00

gda_serialization

Print a better info message for wait_until_finished

2022-06-10 16:57:10 -07:00

jax2tf

Change the remat rule registration to register a default with a GPU specialization.

2022-06-13 21:04:05 -04:00

sparse

[sparse][x64] prevent unnecessary dtype promotion in sparse impls

2022-06-09 14:01:39 -07:00

__init__.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

array.py

Make Shard a normal class making index and replica_id properties on that class. Raise an error if an indices cannot be calculated from a sharding.

2022-06-14 10:37:19 -07:00

callback.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

checkify.py

CI: fix flake8 ignore declarations

2022-04-21 13:44:12 -07:00

global_device_array.py

Add pjit support for Array. Array takes the same codepath as GDA so there are very little modifications to pjit. Add handlers aval, shard_args and result handlers for Array.

2022-06-10 07:32:16 -07:00

host_callback.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

jet.py

correct spelling on comment

2022-06-06 15:24:41 +09:00

loops.py

feat: refactor code using pyupgrade

2022-05-17 22:14:05 +01:00

maps.py

Add Array support to xmap. Just using the GDA path.

2022-06-13 07:31:03 -07:00

mesh_utils.py

[mesh_utils] Avoid relying on process-tiled device order

2022-06-06 13:30:56 -07:00

multihost_utils.py

Move PartitionSpec from sharded_jit.py to pxla.py. The public endpoint is via jax.experimental so that should be used (no changes to the public endpoint).

2022-04-06 15:19:19 -07:00

ode.py

Fix a bug that promoted t to a complex in odeint, and modify a test so it would have caught it

2022-01-07 12:10:35 -05:00

optimizers.py

[JAX] move example libraries from jax.experimental into jax.example_libraries

2021-10-19 17:30:45 -07:00

pjit.py

Allow values in in_axis_resources when Array is an input and make it act like GDA until we remove in_axis_resources. This would also help in transitioning to Array.

2022-06-10 19:55:47 -07:00

sharding.py

Adding jax.Array to jax.experimental. Its pretty much the same as GDA (without the performance optimization for now).

2022-06-06 17:32:00 -07:00

stax.py

[JAX] move example libraries from jax.experimental into jax.example_libraries

2021-10-19 17:30:45 -07:00

x64_context.py

[x64] make x64_context tests work with updated default dtypes

2021-12-08 13:05:40 -08:00