rocm_jax

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

[JAX] Enable C++ device arrays by default.

2021-03-15 10:20:22 -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

Reenable multi-axis all_to_all

2021-03-08 12:45:03 +00:00

callback_test.py

Add scan and while rule for jax.experimental.callback transformation

2021-03-09 19:46:16 -08:00

core_test.py

[avals with names] Add named_shape to ShapedArray and update typecompat

2021-02-05 10:41:05 -08: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

Make jax_debug_nans and jax_debug_infs work with pmap, xmap, and pjit.

2021-03-12 16:22:55 -08:00

djax_test.py

add dynamic-shape-jaxpr experimental file

2021-02-14 13:11:44 -08:00

doubledouble_test.py

Make jax_enable_x64 a thread-local value.

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

dtypes_test.py

BUG: fix jnp.result_type for non-canonical weak types

2021-03-15 14:38:14 -07:00

errors_test.py

mark lax higher-order functions for stack trace filtering

2021-02-24 21:16:20 -08:00

fft_test.py

Make jax_enable_x64 a thread-local value.

2021-02-04 09:48:22 -08: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

Make host_callback_test.py safe to run in environments with > 2 devices.

2021-03-03 18:51:20 +00:00

host_callback_to_tf_test.py

Fix float0 handling

2021-01-06 14:15:42 +02:00

image_test.py

Updating image_test

2020-09-21 16:21:48 -07:00

infeed_test.py

[host_callback] Added host_callback.call to return values to the device

2021-01-04 16:45:48 +02:00

jax_jit_test.py

Increase minimum jaxlib version to 0.1.62.

2021-03-16 15:11:36 -04: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 tolerances in jet_test.py for TPU.

2021-03-03 23:49:20 +00:00

lax_autodiff_test.py

Add workaround for SelectAndScatter padding bug on CPU and GPU.

2021-03-10 15:25:32 -05:00

lax_control_flow_test.py

Increase minimum jaxlib version to 0.1.62.

2021-03-16 15:11:36 -04: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

Make jax_enable_x64 a thread-local value.

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

lax_numpy_test.py

Fix test failure with NumPy 1.20.

2021-03-16 13:46:13 -04:00

lax_numpy_vectorize_test.py

Add ability to specify individual test targets

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

lax_scipy_sparse_test.py

removing x64 enable for testing

2021-02-19 11:17:27 -08:00

lax_scipy_test.py

jax.scipy.special.logsumexp: support integer input

2021-03-04 14:41:30 -08:00

lax_test.py

Increase minimum jaxlib version to 0.1.62.

2021-03-16 15:11:36 -04:00

lax_vmap_test.py

Implement lax.pad batching rule for batched padding values.

2021-02-17 13:30:20 -05:00

linalg_test.py

Remove test that verifies complex eigh is unimplemented on TPU, in preparation for implementing complex eigh on TPU.

2021-03-01 20:29:41 -05:00

loops_test.py

[loops] Extend loops with support for pytrees

2021-01-14 21:17:14 +02:00

masking_test.py

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

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

metadata_test.py

omnistaging, under a flag and disabled by default (#3370 )

2020-07-30 12:59:36 -07:00

multi_device_test.py

Remove the JAX lazy sublanguage.

2021-03-09 21:40:46 -05:00

multibackend_test.py

Add ability to specify individual test targets

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

nn_test.py

Readd omnistaging friendly versions of large constant tests.

2021-03-10 11:34:42 -05: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

Ignore warnings in pjit_test.py

2021-03-05 01:32:36 +00:00

pmap_test.py

Reenable multi-axis all_to_all

2021-03-08 12:45:03 +00:00

polynomial_test.py

Skip unimplemented TPU cases in polynomial_test.py

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

profiler_test.py

fix skip logic for test that requires portpicker

2020-11-08 14:48:14 -08:00

random_test.py

fix the dtype of complex jax.random.normal and add a regression test for it

2021-03-14 22:28:38 +01: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

BUG: fix jax.scipy.stats.dirichlet implementation & add tests

2021-02-17 09:40:18 -08:00

sharded_jit_test.py

Cleanup: remove obsolete jaxlib version checks

2021-02-04 15:13:39 -08:00

stax_test.py

Add ability to specify individual test targets

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

tree_util_test.py

Cleanup: remove obsolete jaxlib version checks

2021-02-04 15:13:39 -08:00

util_test.py

Add ability to specify individual test targets

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

x64_context_test.py

Increase minimum jaxlib version to 0.1.62.

2021-03-16 15:11:36 -04:00

xla_bridge_test.py

fix missing imports

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

xmap_test.py

Increase minimum jaxlib version to 0.1.62.

2021-03-16 15:11:36 -04:00