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 #8191 from mattjj:initial-style-remat

2021-10-12 22:07:55 -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

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -07: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

Adds a Wadler-Lindig pretty printer.

2021-09-27 21:09:24 -04:00

custom_object_test.py

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -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

[JAX] Fix pylint errors.

2021-10-04 17:54:46 -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

Merge pull request #8160 from Jakob-Unfried:main

2021-10-12 09:24:24 -07:00

lax_numpy_einsum_test.py

[JAX] Fix pylint errors.

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

lax_numpy_indexing_test.py

Add DeprecationWarnings to jax.ops.index_... operators.

2021-10-05 20:47:22 -04:00

lax_numpy_test.py

jnp.unique: support size argument with axis

2021-10-12 20:55:27 -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

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -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

Move jax.test_util to jax._src.test_util.

2021-09-24 07:02:49 -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

ahead-of-time lowering and compilation frontend for pjit

2021-10-12 15:01:46 -07:00

pmap_test.py

jnp.array: handle raw device buffers

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

hide keys attribute of PRNGKeyArray in favor of unsafe_raw_array

2021-10-12 08:47:29 -07: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

[sparse] fix padding bug in BCOO._dedupe()

2021-10-11 12:48:27 -07:00

sparsify_test.py

[sparse] add sparsify support for sparse-sparse matmul

2021-10-05 16:45:48 -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

Add --jax_platforms flag to replace --jax_platform_name.

2021-09-30 11:07:52 -07: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