rocm_jax

benchmarks

Cleanup: remove unused imports in tests (#3276 )

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

filecheck

[MLIR] Add a @mlir.cache_lowering decorator that lowers a primitive out-of-line as a reusable function.

2021-12-16 08:34:52 -08: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] Move ann.ann_recall back to tests.

2022-02-24 07:23:17 -08:00

api_test.py

Add separate mechanism for threading name stacks to the lowering

2022-02-23 09:59:09 -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

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22: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

Fix tests and handle cond consts.

2022-02-23 16:11:09 +00:00

compilation_cache_test.py

Run compilation_cache_test.py using TFRT

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

core_test.py

re-enable tests of #8955

2022-01-20 11:11:54 -08:00

custom_linear_solve_test.py

Split lax_control_flow_test into three separate tests.

2022-01-14 13:36:34 -08:00

custom_object_test.py

[MLIR] Change signature of lowering rules.

2021-12-15 19:06:58 -08:00

custom_root_test.py

Split lax_control_flow_test into three separate tests.

2022-01-14 13:36:34 -08:00

debug_nans_test.py

Deprecate maps.mesh and replace it with maps.Mesh.

2022-02-23 10:47:06 -08: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

Run all tests with jax_traceback_filtering=off.

2022-02-15 02:42:58 -08:00

fft_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08: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 ndim and size to GDA

2022-02-10 16:46:26 -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

Deprecate maps.mesh and replace it with maps.Mesh.

2022-02-23 10:47:06 -08: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

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_util_test.py

[JAX] Add support for generating "equation profiles" in JAX.

2022-01-12 14:27:57 -08: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

djax: let make_jaxpr build dyn shape jaxprs

2022-02-01 00:10:21 -08: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

Index update operators: add scatter_apply()

2022-02-18 09:44:40 -08:00

lax_numpy_test.py

BUG: return numpy arrays for jnp.load() with unsupported dtypes

2022-02-25 09:27:42 -08: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

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

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

lax_test.py

Merge pull request #9562 from jakevdp:disable-rank-promotion

2022-02-14 12:27:22 -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

Merge pull request #9544 from SaturdayGenfo:adds-matrix-sqrt

2022-02-17 03:02:47 -08: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 deprecation warning to JaxTestCase and JaxTestLoader

2022-02-17 14:58:58 -08: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

name_stack_test.py

Add separate mechanism for threading name stacks to the lowering

2022-02-23 09:59:09 -08:00

nn_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

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

ode_test.py

Fix a bug that promoted t to a complex in odeint, and modify a test so it would have caught it

2022-01-07 12:10:35 -05: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

Fix the gpu tests that were failing with Future warning

2022-02-23 13:58:28 -08:00

pmap_test.py

Turn execute_replicated into a class so we can access its fields.

2022-02-18 03:18:47 -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

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

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

random_test.py

in tests, compare jnp operations on PRNGKeyArrays to the same on jnp arrays

2022-02-17 11:43:09 -08:00

scipy_fft_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -08: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 some functions for spectral analysis.

2022-02-17 15:59:24 +09:00

scipy_stats_test.py

disable implicit rank promotion in a number of remaining tests

2022-01-28 08:16:30 -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] Updates bcoo_dot_general cuSparse lowering rule by adding sorted indices.

2022-02-14 15:24:14 -08:00

sparsify_test.py

JaxTestCase: set jax_numpy_rank_promotion='raise' by default

2022-02-14 09:22:05 -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

Add deprecation warning to JaxTestCase and JaxTestLoader

2022-02-17 14:58:58 -08:00

transfer_guard_test.py

Implement the JAX transfer guard API

2022-02-14 13:11:49 -08:00

tree_util_test.py

improve pjit in/out_axis_resources pytree errors

2022-02-08 16:23:15 -08:00

util_test.py

Move jax.test_util to jax._src.test_util.

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

Fix the gpu tests that were failing with Future warning

2022-02-23 13:58:28 -08:00