Yash Katariya e32373c3ea Make jnp.array return jax.Array. Add input and result handlers for jax.Array. Also added tests for add under jit. ...

TODO:
* Don't allow `x + y` if `jax.Array` is not fully addressable.
* Figure out how to use the already written tests with Array. Might be able to follow the path taken by SDA.
PiperOrigin-RevId: 457034779

2022-06-24 10:05:06 -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

Simplifies inverting permutation.

2022-06-24 08:06:21 -07:00

sparse

[x64] make sparse tests compatible with strict dtype promotion

2022-06-14 12:45:26 -07:00

__init__.py

CI: fix flake8 ignore declarations

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

array.py

Make jnp.array return jax.Array. Add input and result handlers for jax.Array. Also added tests for add under jit.

2022-06-24 10:05:06 -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

First CL to integrate jax.Array into pmap.

2022-06-16 19:52:31 -07:00

host_callback.py

feat: refactor code using pyupgrade

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

jet.py

[x64] make jet_test compatible with strict dtype promotion

2022-06-21 09:28:24 -07:00

loops.py

feat: refactor code using pyupgrade

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

maps.py

[dynamic-shapes] initial support for dynamic shape typechecks

2022-06-17 14:57:19 -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

[x64] make ode_test pass with strict dtype promotion

2022-06-17 15:21:10 -07:00

optimizers.py

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

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

pjit.py

Merge pull request #11032 from JiahaoYao:patch-1

2022-06-14 16:23:46 -07:00

sharding.py

Check sharding in pmap for jax.Array.

2022-06-22 11:37:10 -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