rocm_jax

benchmarks

Cleanup: remove unused imports in tests (#3276 )

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

filecheck

[MLIR] Keep MLIR IR longer as a Python ir.Module object rather than a string, until it is time to compile it.

2021-12-10 14:56:48 -08:00

notebooks

Update colab CPU test notebook to be more robust

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

third_party/scipy

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

ann_test.py

[JAX] Update JAX users in preparation for a change that makes iteration over a JAX array return JAX arrays, instead of NumPy arrays.

2021-10-28 16:49:37 -07:00

api_test.py

add jax.block_until_ready function

2021-12-14 11:02:14 -08: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

handle trivial case for ppermute batching rule

2021-12-14 10:42:05 -08:00

callback_test.py

Move jax.test_util to jax._src.test_util.

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

checkify_test.py

Add checkify rule for while_loop.

2021-12-08 19:34:20 +00:00

compilation_cache_test.py

[MLIR] Keep MLIR IR longer as a Python ir.Module object rather than a string, until it is time to compile it.

2021-12-10 14:56:48 -08:00

core_test.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

custom_object_test.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08:00

debug_nans_test.py

Increase minimum jaxlib version to 0.1.74.

2021-11-18 15:06:58 -05:00

djax_test.py

[JAX] Fix pylint errors.

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

dtypes_test.py

[x64] add new jax_default_dtype_bits flag

2021-12-07 08:23:37 -08:00

errors_test.py

[JAX] Fix pylint errors.

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

fft_test.py

add fft normalisation

2021-10-20 22:15:35 +01: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

global_device_array_test.py

Add is_fully_replicated API to GDA

2021-12-13 16:18:20 -08: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

[hcb] Simplifications to the host_calback API

2021-12-11 08:24:56 +01: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

Added workaround for bug in XLA

2021-11-18 11:01:50 +02:00

infeed_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

jax_jit_test.py

Move dtype canonicalization out of core.AbstractValue subclasses.

2021-12-07 06:13:07 -08: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

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

lax_autodiff_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

lax_control_flow_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

lax_numpy_einsum_test.py

[JAX] Fix pylint errors.

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

lax_numpy_indexing_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

lax_numpy_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08: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

Add missing gmres test (part of #8849 )

2021-12-08 16:32:33 -08:00

lax_scipy_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

lax_test.py

conv_transpose: allow padding to be specified as a list

2021-11-29 10:31:51 -08:00

lax_vmap_test.py

Implement lax.reduce() batching rule case where batch dimensions differ between operands.

2021-11-03 09:42:01 -04:00

linalg_test.py

Relax test tolerance for slogdet gradient test.

2021-12-07 16:14:06 -05:00

loops_test.py

[x64] make jax.experimental.loops consistent with default dtype

2021-12-08 12:08:49 -08:00

masking_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

mesh_utils_test.py

Add contiguous_submeshes option to mesh_utils.create_device_mesh().

2021-12-10 00:01:12 +00:00

metadata_test.py

[JAX] Fix pylint errors.

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

multi_device_test.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

multibackend_test.py

Cleanup internal representation of XLA translation rules.

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

nn_test.py

fix whitespace

2021-10-27 22:06:06 -07:00

ode_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

optimizers_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

pickle_test.py

Increase minimum jaxlib version to 0.1.74.

2021-11-18 15:06:58 -05:00

pjit_test.py

Handle the edge case where two pjit.FROM_GDA singletons are created.

2021-12-14 01:06:45 +00:00

pmap_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

polynomial_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

profiler_test.py

Fix keyword argument confusion in jax.profiler.annotate_function decorator.

2021-12-06 09:10:44 -05:00

qdwh_test.py

Enable QDWH TPU tests.

2021-11-30 15:47:50 -08:00

random_test.py

Merge pull request #8886 from jakevdp:random-test-x32

2021-12-10 13:35:42 -08: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

stats.multivariate_normal: support broadcasted inputs

2021-10-19 16:58:36 -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

Copybara import of the project:

2021-12-12 09:29:29 -08:00

sparsify_test.py

[sparse] add sparse rule for lax.sub_p

2021-12-13 16:46:47 -08:00

stax_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00

tree_util_test.py

Increase minimum jaxlib version to 0.1.74.

2021-11-18 15:06:58 -05: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

[x64] make x64_context_test more robust

2021-11-22 14:54:30 -08:00

xla_bridge_test.py

Consolidate more XLA-lowering logic between jit, pmap, and xmap.

2021-11-30 14:24:33 -08:00

xla_interpreter_test.py

Split jax.interpreters.xla up into three pieces:

2021-11-22 08:22:43 -08:00

xmap_test.py

Testing: avoid hard-coding random seeds

2021-12-10 10:32:09 -08:00