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

improve error messages for grad(..., has_aux=True)

2021-02-18 09:46:16 -08:00

api_util_test.py

Add ability to specify individual test targets

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

array_interoperability_test.py

Disable some jax2tf tests that fail on GPU.

2021-02-19 15:01:55 -05:00

batching_test.py

Improve error messages for invalid JAX types returned by batched functions.

2021-02-16 20:02:11 -05:00

callback_test.py

Add default implementation of process_custom_jvp_call and

2020-10-16 14:53:23 -07: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

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

2021-02-05 10:41:05 -08:00

debug_nans_test.py

Cleanup: remove obsolete jaxlib version checks

2021-02-04 15:13:39 -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

Re-land #4850 weak types change

2021-02-09 09:07:52 -08:00

errors_test.py

add source line info to more escaped tracer errors

2021-01-18 19:00:04 -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

Disable another test that requires omnistaging

2021-01-14 16:05:14 -08: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

Cleanup: remove obsolete jaxlib version checks

2021-02-04 15:13:39 -08: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

add add_any to jet rules table

2020-12-17 12:10:12 -08:00

lax_autodiff_test.py

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

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

lax_control_flow_test.py

make vmap-of-cond preserve nan safety

2021-02-05 19:12:25 -08:00

lax_numpy_einsum_test.py

Add ability to specify individual test targets

2020-06-29 11:08:57 -07: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

Merge pull request #5731 from terhorst:master

2021-02-22 18:53:38 -08: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

Cleanup: remove remaining instances of rng_factory boilerplate

2020-12-11 13:47:46 -08:00

lax_test.py

Rollback of #5702 due to internal breakage.

2021-02-17 07:32:09 -08: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

Improve error reporting for complex eigh on TPU

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

raise_to_shaped: preserve weak_type by default

2020-10-08 11:53:52 -07:00

multibackend_test.py

Add ability to specify individual test targets

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

nn_test.py

Add support for named axes in jax.nn.one_hot

2021-02-02 15:57:00 +00: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

First draft of pjit.

2021-02-10 01:17:51 +00:00

pmap_test.py

Use the C++ object for the Sharding specification.

2021-02-12 16:02:58 +01:00

polynomial_test.py

Cleanup: remove remaining instances of rng_factory boilerplate

2020-12-11 13:47:46 -08:00

profiler_test.py

fix skip logic for test that requires portpicker

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

random_test.py

Make jax_enable_x64 a thread-local value.

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

Make jax_enable_x64 a thread-local value.

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

xla_bridge_test.py

fix missing imports

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

xmap_test.py

Use the C++ object for the Sharding specification.

2021-02-12 16:02:58 +01:00