rocm_jax

filecheck

(NFC) Prepare for migration from producing MHLO to producing StableHLO

2022-12-15 21:00:07 -08:00

lazy_loader_module

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

notebooks

Update colab CPU test notebook to be more robust

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

third_party/scipy

[x64] more type safety in scipy.optimize.line_search

2022-12-01 14:04:39 -08:00

ann_test.py

[JAX] Disables large k test cases in ann_test.

2022-11-09 11:32:47 -08:00

api_test.py

Check if the sharding input to ShapeDtypeStruct is an instance of Sharding

2023-01-17 12:08:51 -08:00

api_util_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

array_interoperability_test.py

Make JAX tests pass under NumPy 1.24.0rc2.

2022-12-08 19:46:10 +00:00

array_test.py

Bump the minimum jaxlib version to 0.4.1.

2022-12-19 17:49:24 +00:00

batching_test.py

Don't depend on flatten_axis_resources which will error because flatten_axes passes a dummy object() which doesn't work with checks in user pytrees.

2023-01-18 00:13:04 -08:00

BUILD

Make pjit an AxisPrimitive so that it can run the batching rules even if the argument is not batched but there is a axis_index/named shapes inside the pjitted function.

2023-01-18 12:56:07 -08:00

checkify_test.py

Fix scan and name_stack, rewrite cond to use jaxpr_to_checkify_jaxpr.

2023-01-05 16:35:40 +00:00

clear_backends_test.py

Add pjrt_c_api_unimplemented pytest marker to skip unsupported tests.

2023-01-12 22:17:23 +00:00

compilation_cache_test.py

Migrate JAX from producing MHLO to producing StableHLO

2022-12-27 08:53:20 -08:00

core_test.py

Skip core tests which have nested pjits and DShapedArray.

2023-01-13 22:39:31 -08:00

custom_linear_solve_test.py

Fix issue in check_tree, so that custom_linear_solve supports

2022-11-07 21:29:43 +01:00

custom_object_test.py

(NFC) Prepare for migration from producing MHLO to producing StableHLO

2022-12-15 21:00:07 -08:00

custom_root_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

debug_nans_test.py

Remove dependency of maps from pjit to avoid circular imports when importing pjit in api.py.

2022-12-22 13:35:23 -08:00

debugger_test.py

Add pjrt_c_api_unimplemented pytest marker to skip unsupported tests.

2023-01-12 22:17:23 +00:00

debugging_primitives_test.py

Add pjrt_c_api_unimplemented pytest marker to skip unsupported tests.

2023-01-12 22:17:23 +00:00

dtypes_test.py

Don't run FP8 dtype test on TPU.

2023-01-18 11:22:17 -08:00

dynamic_api_test.py

(NFC) Prepare for migration from producing MHLO to producing StableHLO

2022-12-15 21:00:07 -08:00

errors_test.py

Add changes accidentally omitted from

2022-10-10 19:11:58 +00:00

experimental_rnn_test.py

[ROCm]: Disable some tests on ROCm platform

2022-12-19 21:33:13 +00:00

fft_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

for_loop_test.py

[x64] minor test-only updates for better type safety

2022-11-30 15:18:40 -08:00

generated_fun_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

gfile_cache_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

global_device_array_test.py

Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL).

2022-11-14 14:44:00 -08:00

heap_profiler_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

host_callback_test.py

[host_callback] Add device_index to hcb.call and add tests

2023-01-18 12:41:11 -08:00

host_callback_to_tf_test.py

Add pjrt_c_api_unimplemented pytest marker to skip unsupported tests.

2023-01-12 22:17:23 +00:00

image_test.py

Add changes accidentally omitted from

2022-10-10 19:11:58 +00:00

infeed_test.py

Add pjrt_c_api_unimplemented pytest marker to skip unsupported tests.

2023-01-12 22:17:23 +00:00

jax_jit_test.py

Error on numpy masked array inputs.

2022-12-27 15:42:49 -08:00

jax_to_ir_test.py

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08:00

jaxpr_effects_test.py

Add in effects_barrier for the pmap unordered callback test

2023-01-16 14:44:44 -08:00

jaxpr_util_test.py

Make jaxpr_util_test work with jit/pjit merge

2023-01-09 16:50:04 -08:00

jet_test.py

relax tanh test tols for upcoming xla change

2022-12-20 21:06:09 -08:00

lax_autodiff_test.py

relax tanh test tols for upcoming xla change

2022-12-20 21:06:09 -08:00

lax_control_flow_test.py

Fix excessive recompiles in lax.cond

2023-01-18 10:17:01 -08:00

lax_numpy_einsum_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lax_numpy_indexing_test.py

[x64] more type safety for lax_numpy-related tests

2022-12-01 11:18:02 -08:00

lax_numpy_operators_test.py

Re-land https://github.com/google/jax/pull/11996

2023-01-09 16:57:32 -08:00

lax_numpy_reducers_test.py

Make JAX tests pass under NumPy 1.24.0rc2.

2022-12-08 19:46:10 +00:00

lax_numpy_test.py

mgrid: Fix zero-length meshgrid

2023-01-12 19:43:28 +01:00

lax_numpy_vectorize_test.py

Port more tests from jtu.cases_from_list to jtu.sample_product.

2022-10-11 21:06:08 +00:00

lax_scipy_sparse_test.py

Migrate remaining tests from jtu.cases_from_list to jtu.sample_product.

2022-10-12 15:20:53 +00:00

lax_scipy_test.py

[x64] make tests more type-safe

2022-12-02 13:21:35 -08:00

lax_test.py

Fix scatter in CLIP mode with uint32 and uint64 indices

2023-01-17 06:26:21 -08:00

lax_vmap_test.py

Reenable GPU int matmul test since the XLA bug is fixed.

2022-12-13 13:50:42 -08:00

lazy_loader_test.py

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

linalg_test.py

[ROCm]: Disable some tests on ROCm platform

2022-12-19 21:33:13 +00:00

lobpcg_test.py

Switch to new-style f-strings

2022-12-01 09:14:16 -08:00

mesh_utils_test.py

Allow more cases in _TRANSPOSE_TRICKS by ignoring leading 1s in the mesh shape.

2022-12-28 09:52:24 -08:00

metadata_test.py

Delete jax_experimental_name_stack flag

2022-11-10 11:59:50 -08:00

monitoring_test.py

Add JAX monitoring library that instruments code via events.

2022-11-15 12:41:41 -08:00

multi_device_test.py

Fix the error message of different devices when jit/pjit are merged

2023-01-09 09:03:55 -08:00

multibackend_test.py

pjit allows nesting of pjits where the outer backend is None while the inner backend is something other than device_under_test(). This is because the inner backend will take priority.

2023-01-17 16:39:45 -08:00

multiprocess_gpu_test.py

Migrate pytestmark usage to new @jtu.pytest_mark_if_available decorator.

2023-01-12 22:44:39 +00:00

name_stack_test.py

jit was the default name in name_stack in mlir.py. Fix that by taking the name as an optional argument (defaulting to jit) so that nested pjits will show up as pjit in the name stack.

2023-01-13 15:00:22 -08:00

nn_test.py

[x64] minor test-only updates for better type safety

2022-11-30 15:18:40 -08:00

ode_test.py

Merge pull request #11077 from DanPuzzuoli:ode_dt_max

2022-10-28 16:05:25 -07:00

optimizers_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

pickle_test.py

Bump the minimum jaxlib version to 0.4.1.

2022-12-19 17:49:24 +00:00

pjit_test.py

Add support for __signature__ to PjitFunction.

2023-01-17 17:28:14 -08:00

pmap_test.py

Migrate pytestmark usage to new @jtu.pytest_mark_if_available decorator.

2023-01-12 22:44:39 +00:00

polynomial_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

profiler_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

python_callback_test.py

Add IO callback

2023-01-17 13:55:05 -08:00

qdwh_test.py

[x64] minor test-only updates for better type safety

2022-11-30 15:18:40 -08:00

random_test.py

[ROCm]: Disable some tests on ROCm platform

2022-12-19 21:33:13 +00:00

remote_transfer_test.py

Migrate pytestmark usage to new @jtu.pytest_mark_if_available decorator.

2023-01-12 22:44:39 +00:00

scipy_fft_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

scipy_interpolate_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

scipy_ndimage_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

scipy_optimize_test.py

Port more tests from jtu.cases_from_list to jtu.sample_product.

2022-10-11 21:06:08 +00:00

scipy_signal_test.py

Rename test flag --num_generated_cases to --jax_num_generated_cases.

2022-11-09 18:58:05 -08:00

scipy_stats_test.py

Added scipy.stats.bernoulli cdf and ppf.

2022-12-22 18:12:25 +00:00

sparse_test.py

[sparse] add higher-level version of bcoo_extract & improve tests

2023-01-13 07:13:13 -08:00

sparsify_test.py

[sparse] avoid re-indexing for linear unary ops

2022-11-17 16:31:46 -08:00

stack_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

state_test.py

Add support for Ellipsis as an index for stateful operations.

2022-12-23 22:46:50 -08:00

stax_test.py

[x64] more type safety in stax_test.py

2022-12-02 10:02:25 -08:00

svd_test.py

[x64] make tests more type-safe

2022-12-02 13:21:35 -08:00

transfer_guard_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

tree_util_test.py

Store sorted flattened dict keys in PyTree as a c++ vector instead of py::list to avoid creating new python object on every single dict flatten. For deeply nested dict, this avoids excessive gc pressure and avoids the slowdown whenever gc needs to sweep too many live python objects.

2023-01-18 13:40:43 -08:00

typing_test.py

[x64] minor test-only updates for better type safety

2022-11-30 15:18:40 -08:00

util_test.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

version_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

x64_context_test.py

Copybara import of the project:

2022-10-10 11:35:32 -07:00

xla_bridge_test.py

Fix xla_bridge_test on TPU

2022-10-17 15:44:06 -07:00

xla_interpreter_test.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

xmap_test.py

Migrate pytestmark usage to new @jtu.pytest_mark_if_available decorator.

2023-01-12 22:44:39 +00:00