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

2022-06-10 00:17:13 +00: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

Add a Async checkpoint manager interface.

2022-06-08 08:33:41 -07:00

jax2tf

[jax2tf] Rewrites the converters_eval framework and makes various improvements.

2022-06-02 22:39:09 -07:00

sparse

[sparse] Track unique_indices in sparse transform.

2022-06-03 14:51:26 -07:00

__init__.py

CI: fix flake8 ignore declarations

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

array.py

Add the Shard class from GDA to jax.Array to make the management of shards easier and less verbose.

2022-06-06 18:45:15 -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 the Shard class from GDA to jax.Array to make the management of shards easier and less verbose.

2022-06-06 18:45:15 -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

Enable debug print in xmap

2022-06-07 14:05:46 -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

Support a mix of user sharding and auto sharding in pjit.

2022-06-06 09:27:17 -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