rocm_jax

benchmarks

Cleanup: remove unused imports in tests (#3276 )

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

notebooks

Update colab CPU test notebook to be more robust

2021-07-13 09:30:56 -07:00

third_party/scipy

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

api_test.py

Merge pull request #8242 from mattjj:remat-fix6

2021-10-15 17:31:06 -07:00

api_util_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

array_interoperability_test.py

update test of dlpack error message

2021-10-07 19:14:13 -07:00

batching_test.py

Enable batcher and batched collective rules for tiled all gathers

2021-10-15 14:37:38 +00:00

callback_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

compilation_cache_test.py

Include XLA_FLAGS in persistent compilation cache key.

2021-10-06 14:11:40 -07:00

core_test.py

Change jax.core.DropVar to be a non-singleton.

2021-10-18 15:02:54 -07:00

custom_object_test.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

debug_nans_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

djax_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

dtypes_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

errors_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

fft_test.py

[JAX] Fix more pylint errors.

2021-10-04 16:37:15 -07:00

file_system_cache_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

generated_fun_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

heap_profiler_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

host_callback_test.py

Change jax.core.DropVar to be a non-singleton.

2021-10-18 15:02:54 -07:00

host_callback_to_tf_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

image_test.py

[JAX] Canonicalize shapes in jax.image.resize().

2021-09-30 12:38:39 -07:00

infeed_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jax_jit_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jax_to_hlo_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jaxpr_util_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

jet_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

lax_autodiff_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

lax_control_flow_test.py

Consolidate primitive and jit lowering paths.

2021-10-13 12:36:53 -04:00

lax_numpy_einsum_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

lax_numpy_indexing_test.py

Add a regression test for #7461 .

2021-10-13 11:11:24 -04:00

lax_numpy_test.py

jnp.bincount: fix corner cases

2021-10-15 12:31:17 -07:00

lax_numpy_vectorize_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

lax_scipy_sparse_test.py

Reenable a test that was disabled due to an (apparently fixed) LLVM bug.

2021-10-16 10:34:36 -07:00

lax_scipy_test.py

Relax the tolerances in one 'lax_scipy_test' case

2021-09-28 17:53:51 -07:00

lax_test.py

fix rng_bit_generator translation rule

2021-10-05 20:19:57 -07:00

lax_vmap_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

linalg_test.py

schur lapack wrapper

2021-09-29 14:29:52 +02:00

loops_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

masking_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

metadata_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

multi_device_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

multibackend_test.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

nn_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

ode_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

optimizers_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

pickle_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

pjit_test.py

Cleanup internal representation of XLA translation rules.

2021-10-16 07:53:24 -07:00

pmap_test.py

add more grad-of-jit/pmap caching tests

2021-10-13 11:06:17 -07:00

polynomial_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -07:00

profiler_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

random_test.py

Cleanup random.permutation

2021-10-13 14:13:00 -06:00

scipy_fft_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_ndimage_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_optimize_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_signal_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

scipy_stats_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

sharded_jit_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

sparse_test.py

Merge pull request #8269 from jakevdp:sparse-dtype

2021-10-19 05:45:24 -07:00

sparsify_test.py

[sparse] respect weak types in sparsify transform

2021-10-18 13:35:20 -07:00

stax_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

tree_util_test.py

Fixes for Python 3.10.

2021-10-05 15:25:28 -04:00

util_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

x64_context_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

xla_bridge_test.py

Remove xla_bridge.make_computation_builder().

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

xla_interpreter_test.py

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07:00

xmap_test.py

Use shaped_abstractify to allow abstract arguments in .lower methods

2021-09-27 05:32:12 -07:00