rocm_jax

benchmarks

Cleanup: remove unused imports in tests (#3276 )

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

filecheck

Remove the jax_enable_mlir flag. MLIR is now the only supported code path.

2022-04-04 08:40:09 -07:00

notebooks

Update colab CPU test notebook to be more robust

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

third_party/scipy

Fix test failure in line search test due to scipy 1.8 release.

2022-02-07 14:49:54 +00:00

ann_test.py

[JAX] Fix batch logic for approx_min/max_k

2022-04-08 13:50:36 -07:00

api_test.py

Drop dead code now that the minimum jaxlib version is 0.3.2.

2022-04-13 13:34:00 -04: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

vmap: preserve weak_type in batching tracer

2022-03-30 11:06:56 -07:00

callback_test.py

Move jax.test_util to jax._src.test_util.

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

checkify_test.py

Checkify: add axis and axis size to OOB error message.

2022-03-31 15:16:35 +01:00

compilation_cache_test.py

Run compilation_cache_test.py using TFRT

2022-02-04 10:17:59 -08:00

core_test.py

prototyping dynamic shapes

2022-04-11 22:10:47 -07:00

custom_linear_solve_test.py

Update JAX to use new math libraries in ROCm-5.0.

2022-03-01 20:02:15 +00:00

custom_object_test.py

[MHLO] Add direct MHLO lowerings for sparse primitives.

2022-04-08 08:43:57 -07:00

custom_root_test.py

Relax tolerances slightly for MKL.

2022-02-25 22:02:55 +00:00

debug_nans_test.py

debug_nans: don't return results of successfully running de-optimized function

2022-04-12 14:40:19 -07:00

dtypes_test.py

remove _convert_element_type from public jax.lax module

2022-03-09 18:46:38 -08:00

errors_test.py

Increase minimum jaxlib version to 0.3.0.

2022-03-04 10:33:03 -05:00

fft_test.py

Skip test outside x64

2022-04-04 16:00:18 -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

global_device_array_test.py

Make mesh_axes on GDA strict by only allowing PartitionSpecs to be consistent with pjit.

2022-03-07 08:59:23 -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

[MHLO] Change jax.xla_computation() to use MHLO lowering internally.

2022-04-13 06:28:38 -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

Make negative dimension return consistent results for image.scale_and_translate

2022-02-24 22:34:00 -05:00

infeed_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

jax_jit_test.py

jax_jit_test: inherit from JaxTestCase

2022-01-24 15:32:59 -08:00

jax_to_ir_test.py

Add build time support for AOT compilation to TF graphs.

2021-12-20 06:18:39 -08:00

jaxpr_effects_test.py

Add tests for higher order primitives

2022-04-12 18:12:44 -07:00

jaxpr_util_test.py

Increase minimum jaxlib version to 0.3.0.

2022-03-04 10:33:03 -05:00

jet_test.py

Add test and jet-primitive for dynamic_slice

2022-02-08 13:28:41 +01:00

lax_autodiff_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

lax_control_flow_test.py

better errors: check for callability of lax.control_flow arguments

2022-04-13 10:39:01 -07:00

lax_numpy_einsum_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

lax_numpy_indexing_test.py

segment_max: fix identity for boolean dtype

2022-03-15 09:20:20 -07:00

lax_numpy_test.py

Fix polydiv kokoro tests

2022-04-13 13:21:29 -07:00

lax_numpy_vectorize_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

lax_scipy_sparse_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

lax_scipy_test.py

Update JAX to use new math libraries in ROCm-5.0.

2022-03-01 20:02:15 +00:00

lax_test.py

Adds simple effect types to jaxprs

2022-04-11 11:50:41 -07:00

lax_vmap_test.py

fixup from 5415306: remove extraneous lines

2022-03-11 15:19:10 -08:00

linalg_test.py

Add jax.scipy.linalg.rsf2csf

2022-04-06 21:06:23 +01: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

Allow a single logical mesh dim to take all devices.

2022-03-16 20:58:59 -07:00

metadata_test.py

[MHLO] Change jax.xla_computation() to use MHLO lowering internally.

2022-04-13 06:28:38 -07:00

multi_device_test.py

Update JAX to use new math libraries in ROCm-5.0.

2022-03-01 20:02:15 +00:00

multibackend_test.py

Cleanup internal representation of XLA translation rules.

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

name_stack_test.py

[MHLO] Change jax.xla_computation() to use MHLO lowering internally.

2022-04-13 06:28:38 -07:00

nn_test.py

Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning.

2022-03-23 20:55:22 +00:00

ode_test.py

[XLA:CPU] Relax test tolerances for tests using XLA:CPU.

2022-02-28 09:26:54 -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

Drop dead code now that the minimum jaxlib version is 0.3.2.

2022-04-13 13:34:00 -04:00

pmap_test.py

Deprecate jax.tree_util.tree_multimap

2022-04-01 14:51:54 -07: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

Update JAX to use new math libraries in ROCm-5.0.

2022-03-01 20:02:15 +00:00

random_test.py

2022-04-08 01:08:55 -07:00

remote_transfer_test.py

Disable remote_transfer_test on GPU. It currently crashes.

2022-04-14 08:06:27 -07:00

scipy_fft_test.py

Update JAX to use new math libraries in ROCm-5.0.

2022-03-01 20:02:15 +00:00

scipy_interpolate_test.py

Merge pull request #8625 from Edenhofer:regular_grid_interpolator

2022-01-24 12:02:34 -08:00

scipy_ndimage_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

scipy_optimize_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

scipy_signal_test.py

Add istft to jax.scipy.signal.

2022-04-01 14:28:53 +09:00

scipy_stats_test.py

Fix#10219

2022-04-13 04:04:11 +08:00

sharded_jit_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08:00

sparse_test.py

[sparse] Use sorted indices instead of sorted rows only.

2022-04-09 08:33:48 -07:00

sparsify_test.py

[sparse] refactor internal implementation of sparsify transform

2022-03-07 12:48:03 -08:00

stax_test.py

disable implicit rank promotion in a number of remaining tests

2022-01-28 08:16:30 -08:00

svd_test.py

Update JAX to use new math libraries in ROCm-5.0.

2022-03-01 20:02:15 +00:00

transfer_guard_test.py

Bump minimum jaxlib version to 0.3.2 and remove transfer guard compatibility code

2022-04-11 15:33:27 +00:00

tree_util_test.py

Deprecate jax.tree_util.tree_multimap

2022-04-01 14:51:54 -07:00

util_test.py

Fix race condition for weakref destructor by catching rare exceptions.

2022-04-01 12:04:36 -07:00

version_test.py

Introduce a new jax/jaxlib versioning scheme.

2022-02-07 17:59:42 -05: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

[MHLO] Change jax.xla_computation() to use MHLO lowering internally.

2022-04-13 06:28:38 -07:00