Skye Wanderman-Milne 9128ba0c74 Replace host_id with process_index terminology, take 2. ...

We're switching to the new terminology to avoid confusion in cases
where multiple jax processes are running on a single host, and each
process has a unique process_index/host_id.

This keeps aliases for the old `host_id` APIs for now, but these will
eventually be removed.

This was originally commited in
b77ef5138b631378e6a8ceb8bafc94fe91239bae, but reverted in
14acd070c2afb11c81fc91f43790577cd48cbf67 due to Google-internal test
failures from renaming the local_devices argument name. This change is
identical except it also adds staging for the argument name change.

2021-04-20 18:13:34 -07:00

..

benchmarks

Cleanup: remove unused imports in tests (#3276)

2020-06-01 11:49:35 -07:00

notebooks

Add Colab test notebooks for CPU, GPU, and TPU (#3000)

2020-05-08 11:11:42 -07:00

third_party/scipy

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

api_test.py

Merge pull request #6239 from j-towns:lt-allow-integers

2021-04-20 10:23:10 -07:00

api_util_test.py

Add ability to specify individual test targets

2020-06-29 11:08:57 -07:00

array_interoperability_test.py

Add a test that fails due to unsupported strides.

2021-03-10 14:19:21 -08:00

batching_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

callback_test.py

Add support for custom derivatives in jax.experimental.callback

2021-03-22 12:57:33 -07:00

core_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

custom_object_test.py

[JAX] Add a .weak_type attribute to C++ array objects.

2021-03-15 12:30:15 -07:00

debug_nans_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

djax_test.py

support implicit broadcasting in transpose rules

2021-04-16 12:51:11 -07:00

doubledouble_test.py

Make jax_enable_x64 a thread-local value.

2021-02-04 09:48:22 -08:00

dtypes_test.py

[JAX] Move contents of jax/dtypes.py to jax/_src/dtypes.py.

2021-04-07 19:35:51 -07:00

errors_test.py

errors: add NonConcreteBooleanIndexError & debugging tips

2021-03-23 11:23:20 -07:00

fft_test.py

Add support for padding and cropping to fft

2021-04-17 08:38:24 +02:00

generated_fun_test.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

host_callback_test.py

xla: improvement to string representation of PyTreeDef

2021-04-19 14:06:36 -07:00

host_callback_to_tf_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

image_test.py

Move jax.config to jax._src.config.

2021-04-19 08:53:12 -07:00

infeed_test.py

Move jax.config to jax._src.config.

2021-04-19 08:53:12 -07:00

jax_jit_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

jax_to_hlo_test.py

Small cleanups to dependency structure.

2021-01-20 12:43:28 -08:00

jaxpr_util_test.py

Make jax_enable_x64 a thread-local value.

2021-02-04 09:48:22 -08:00

jet_test.py

Adjust JetTest.test_scatter_add tolerance for TPU

2021-03-31 17:25:27 +00:00

lax_autodiff_test.py

support implicit broadcasting in transpose rules

2021-04-16 12:51:11 -07:00

lax_control_flow_test.py

xla: improvement to string representation of PyTreeDef

2021-04-19 14:06:36 -07:00

lax_numpy_einsum_test.py

Try to avoid transposes in jnp.einsum by considering both argument orders to dot_general.

2021-03-04 10:00:08 -05:00

lax_numpy_indexing_test.py

Merge pull request #6420 from apaszke:faster-tests

2021-04-14 15:24:06 -07:00

lax_numpy_test.py

Make jnp.unique() support complex inputs

2021-04-20 13:57:30 -07:00

lax_numpy_vectorize_test.py

Mark arguments to jax.jit() other than the function as keyword-only.

2021-04-08 10:32:35 -04:00

lax_scipy_sparse_test.py

Disable lax_scipy_sparse_test.py cases that are hanging on GPU.

2021-04-15 15:54:15 -07:00

lax_scipy_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

lax_test.py

Merge pull request #6420 from apaszke:faster-tests

2021-04-14 15:24:06 -07:00

lax_vmap_test.py

Merge pull request #6420 from apaszke:faster-tests

2021-04-14 15:24:06 -07:00

linalg_test.py

Reimplement as JAX Primitive

2021-04-03 14:11:36 -07:00

loops_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

masking_test.py

Merge pull request #6375 from gnecula:mask_clean

2021-04-12 05:50:19 -07:00

metadata_test.py

remove disable_omnistaging mechanism

2021-03-29 15:26:57 -07:00

multi_device_test.py

remove disable_omnistaging mechanism

2021-03-29 15:26:57 -07:00

multibackend_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

nn_test.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

ode_test.py

Fix scan carry types in gradient of complex ODE (#4130)

2020-08-24 13:50:44 -07:00

optimizers_test.py

X32 tests: fail on dtype warnings

2020-12-08 13:03:30 -08:00

pjit_test.py

Extend pjit error checking to rank errors too

2021-04-20 11:39:58 -07:00

pmap_test.py

Speed up test generation 7x

2021-04-14 15:58:05 +00:00

polynomial_test.py

Skip unimplemented TPU cases in polynomial_test.py

2021-03-04 00:12:55 +00:00

profiler_test.py

Merge pull request #6380 from skye:python_trace_test

2021-04-14 12:29:16 -07:00

random_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

scipy_ndimage_test.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

scipy_optimize_test.py

Enforce the the args argument to jax.scipy.optimize.minimize is a tuple.

2021-02-16 15:57:46 -05:00

scipy_signal_test.py

Relax some test tolerances of tests that fail on GPU in -x32 mode.

2021-02-16 22:14:43 -05:00

scipy_stats_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

sharded_jit_test.py

Enable sharded_jit_test and xmap SPMD tests for GPU.

2021-03-23 10:22:10 -07:00

sparse_ops_test.py

Add experimental/sparse_ops & cusparse wrappers in jaxlib

2021-04-15 10:11:10 -07:00

stax_test.py

Add ability to specify individual test targets

2020-06-29 11:08:57 -07:00

tree_util_test.py

Use assertRegex in TreeTest since pytest expects a different module prefix.

2021-04-20 14:10:30 -07:00

util_test.py

Add ability to specify individual test targets

2020-06-29 11:08:57 -07:00

x64_context_test.py

Move jax.api to jax._src.api.

2021-04-13 09:43:24 -07:00

xla_bridge_test.py

Replace host_id with process_index terminology, take 2.

2021-04-20 18:13:34 -07:00

xmap_test.py

Fix cache misses when re-creating equivalent mesh objects

2021-04-20 03:48:26 -07:00