rocm_jax

filecheck

Migrate custom_call filecheck to use internal custom_call since the external one is deprecated.

2025-04-07 17:32:02 -07:00

mosaic

Fix Pallas tests so they work with JAX_TEST_NUM_THREADS >= 1.

2025-04-10 16:57:34 -07:00

notebooks

…

pallas

[Pallas] Fix integer array indexing

2025-04-10 19:10:35 -07:00

testdata

…

third_party/scipy

…

ann_test.py

…

aot_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

api_test.py

Add JVP rule for linear_call.

2025-04-10 09:12:01 -04:00

api_util_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

array_api_skips.txt

[array_api] update array_api_version to 2024.12

2025-04-02 13:46:07 -07:00

array_api_test.py

[array api] return all devices in devices()

2025-03-31 08:50:39 -07:00

array_extensibility_test.py

jnp.repeat: don't cast repeats to array, as they must be static.

2025-04-08 10:32:03 -07:00

array_interoperability_test.py

[dlpack] Support more DLPack dtypes now that we target DLPack 1.1

2025-03-14 09:10:56 -07:00

array_test.py

Raise an error if the type passed to axis_types argument of Mesh and AbstractMesh is not jax.sharding.AxisType.

2025-04-06 23:38:09 -07:00

attrs_test.py

PR #27576 : [attrs] experimental appendattr

2025-03-28 15:21:12 -07:00

batching_test.py

[vmappable] fix trace context bugs

2025-04-06 00:33:40 +00:00

blocked_sampler_test.py

Make blocked_fold_in consistent when the block sizes induce padding

2025-03-18 09:12:11 -07:00

BUILD

Update test shardings.

2025-04-10 14:01:15 -07:00

cache_key_test.py

harden cache against jaxlib ver

2025-04-07 23:30:31 +00:00

checkify_test.py

optimize while_loop by moving readonly carry components to be consts

2025-04-11 00:48:52 +00:00

cholesky_update_test.py

…

clear_backends_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

colocated_python_test.py

Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays.

2025-03-06 14:57:52 -08:00

compilation_cache_test.py

#sdy unskip JAX Shardy tests that are already passing

2025-01-23 09:26:38 -08:00

config_test.py

…

core_test.py

Mark as thread-unsafe tests that modify possibly-cached jaxprs in-place.

2025-04-10 11:40:42 -07:00

cudnn_fusion_test.py

…

custom_linear_solve_test.py

Handle effects in lax.custom_linear_solve.

2025-02-03 11:14:48 -05:00

custom_partitioning_sharding_rule_test.py

[jax:custom_partitioning] Support SdyShardingRule with multiple leading

2025-01-23 08:20:38 -08:00

custom_root_test.py

…

debug_info_test.py

Fix linear_call to allow recursive definitions.

2025-04-09 09:45:32 -04:00

debug_nans_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

debugger_test.py

[JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly.

2025-01-08 06:37:44 -08:00

debugging_primitives_test.py

Add option for debug print to be called on partitioned arguments rather than all-gather and print the entire logical arguments.

2025-04-09 12:00:00 -07:00

deprecation_test.py

[internal] add deprecation test utilities

2025-01-10 11:54:09 -08:00

device_test.py

…

distributed_initialize_test.py

Fixed broken JAX distributed tests.

2025-03-25 09:39:10 -07:00

distributed_test.py

Fixed broken JAX distributed tests.

2025-03-25 09:39:10 -07:00

dtypes_test.py

cleanup now that we depend on ml_dtypes>=0.5

2025-03-28 07:44:38 -07:00

dynamic_api_test.py

Port tests away from setUpClass and setUpModule to setUp alone.

2025-01-08 08:14:50 -08:00

error_check_test.py

Prototype of adding error checking to jax.numpy functions

2025-03-26 10:37:34 -07:00

errors_test.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

experimental_rnn_test.py

Fix integer overflow in workspace size computations for experimental.rnn.*.

2025-03-12 08:22:04 -07:00

export_back_compat_test.py

[export] Add backwards compatibility test for annotate_device_placement.

2025-04-08 00:10:23 -07:00

export_harnesses_multi_platform_test.py

[export] Fixes for export_harnesses_multi_platform_test.

2025-01-29 06:12:33 -08:00

export_test.py

fix bug in export_module when no mesh axes are empty for shardy.

2025-04-10 09:21:58 -07:00

extend_test.py

[better_errors] Make it explicit that debug_info is not None.

2025-02-13 22:07:04 -08:00

ffi_test.py

Finalize deprecation of ffi_call with inline arguments.

2025-04-08 13:09:42 -07:00

fft_test.py

…

for_loop_test.py

…

fused_attention_stablehlo_test.py

Removed unused import

2025-03-05 15:18:43 +00:00

garbage_collection_guard_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

generated_fun_test.py

Add a variant of safe_map() that has no return value, named foreach().

2025-03-14 07:42:48 -07:00

gpu_memory_flags_test.py

…

heap_profiler_test.py

…

image_test.py

…

infeed_test.py

Make JAX test suite pass (at least most of the time) with multiple threads enabled.

2025-01-10 06:58:46 -08:00

jax_jit_test.py

Add test case for jit compile of vmap on gpu.

2025-01-23 17:50:19 -08:00

jax_numpy_error_test.py

Temporarily skip JaxNumpyErrorTests in multi-thread environments

2025-04-07 07:22:17 -07:00

jax_to_ir_test.py

cleanup now that we depend on ml_dtypes>=0.5

2025-03-28 07:44:38 -07:00

jaxpr_effects_test.py

Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.

2025-02-21 09:45:59 -08:00

jaxpr_util_test.py

…

jet_test.py

…

key_reuse_test.py

[key reuse] fix signature for device_put

2025-01-17 09:47:50 -08:00

lax_autodiff_test.py

Fix lax_autodiff_test on v5p

2025-03-19 15:41:28 -07:00

lax_control_flow_test.py

optimize while_loop by moving readonly carry components to be consts

2025-04-11 00:48:52 +00:00

lax_metal_test.py

…

lax_numpy_einsum_test.py

…

lax_numpy_indexing_test.py

_attempt_rewriting_take_via_slice(): canonicalize the slice index before checking it's not too long, so that e.g. my_1d_array[:, ...] can be treated as a slice rather than generating a gather operation.

2025-04-04 10:10:38 -07:00

lax_numpy_operators_test.py

…

lax_numpy_reducers_test.py

Skip failing tests on TPU v6+

2025-03-28 07:17:39 -07:00

lax_numpy_test.py

Merge tuple_replace and tuple_update in jax._src.util.

2025-04-09 12:50:42 -04:00

lax_numpy_ufuncs_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

lax_numpy_vectorize_test.py

…

lax_scipy_sparse_test.py

…

lax_scipy_special_functions_test.py

Fix betainc edge cases and inaccuracies when a is close to zero.

2025-03-20 00:54:55 +02:00

lax_scipy_spectral_dac_test.py

…

lax_scipy_test.py

Skip failing tests on TPU v6+

2025-03-28 07:17:39 -07:00

lax_test.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

lax_vmap_op_test.py

…

lax_vmap_test.py

…

layout_test.py

Add **experimental** with_dll_constraint API. This is for cases when the users wants to let SPMD decide the sharding.

2025-04-07 16:21:58 -07:00

lazy_loader_test.py

…

linalg_sharding_test.py

Temporarily skip some more linalg sharding checks.

2025-02-28 12:26:18 -08:00

linalg_test.py

jnp.linalg: add symmetrize_input argument & docs

2025-04-07 14:46:38 -07:00

lobpcg_test.py

Disabled default env var LOBPCG_EMIT_DEBUG_PLOTS=1

2025-04-02 15:12:12 +00:00

logging_test.py

Fix logging_test fails on Linux with NVIDIA Driver only.

2025-03-18 14:51:35 -07:00

lru_cache_test.py

Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1

2025-02-14 12:01:55 -08:00

magma_linalg_test.py

Remove code present to support jaxlib < 0.5.1.

2025-02-26 07:40:40 -05:00

memories_test.py

Add a 2D test in memories_test.

2025-04-10 21:32:56 -07:00

mesh_utils_test.py

Add license headers to files that were missing them.

2025-03-27 08:45:15 -07:00

metadata_test.py

…

mock_gpu_test.py

Annotate several tests as thread-unsafe.

2025-01-10 11:24:39 -08:00

mock_gpu_topology_test.py

Annotate several tests as thread-unsafe.

2025-01-10 11:24:39 -08:00

monitoring_test.py

Add scalar event logging function

2025-03-28 23:23:42 +00:00

mosaic_test.py

…

multi_device_test.py

Make JAX test suite pass (at least most of the time) with multiple threads enabled.

2025-01-10 06:58:46 -08:00

multibackend_test.py

…

multiprocess_gpu_test.py

add distributed.is_initialized

2025-02-18 16:47:19 -08:00

mutable_array_test.py

[mutable-arrays] limit implicit ref_swap dtype promotion

2025-04-08 10:25:10 -07:00

name_stack_test.py

Just use jit as the string in error messages instead of jit and pjit based on resource_env. This is to start deprecating the need for with mesh and replace it with use_mesh(mesh).

2025-03-05 20:09:30 -08:00

nn_test.py

Enable public doc for scaled dot

2025-03-27 00:05:28 +00:00

ode_test.py

[XLA:CPU] Use central difference to calculate numerical gradient

2025-01-22 07:49:43 -08:00

optimizers_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

package_structure_test.py

…

pgle_test.py

AutoPGLE: force-disable graphs less

2025-03-31 18:01:56 +00:00

pickle_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

pjit_test.py

Add jaxlib_extension_version guard against explicit copying

2025-04-07 13:50:15 -07:00

pmap_test.py

Removed eager_pmap config option

2025-04-08 03:30:36 -07:00

polynomial_test.py

…

pretty_printer_test.py

…

profiler_test.py

Add a profiler test for gpu run

2025-02-28 13:45:46 -08:00

python_callback_test.py

Reverts 735cec18cb2f8dff2aea5e503fd886a37aee094e

2025-04-07 07:40:14 -07:00

pytorch_interoperability_test.py

…

qdwh_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

ragged_collective_test.py

Add simple vmap support for lax.ragged_all_to_all.

2025-04-02 12:10:34 -07:00

random_lax_test.py

Refactor random_lax_test.py

2025-04-09 15:33:04 -07:00

random_test.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

roofline_test.py

Support conv unfused_flops in roofline.

2025-04-02 14:08:30 -07:00

scaled_matmul_stablehlo_test.py

Improve based on review 1

2025-03-20 23:26:21 +00:00

scipy_fft_test.py

…

scipy_interpolate_test.py

…

scipy_ndimage_test.py

…

scipy_optimize_test.py

…

scipy_signal_test.py

…

scipy_spatial_test.py

test: change random matrix generation for Rotation

2025-01-31 09:21:46 -08:00

scipy_stats_test.py

feat(gh-13291): Add exponential distribution functions: cdf, logcdf, sf, logsf, and ppf

2025-02-01 12:51:11 +05:00

shape_poly_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

shard_alike_test.py

Use the mesh of out_aval when converting GSPMDSharding to NamedSharding. This makes sure that the axis types of the corresponding output is correct.

2025-02-20 17:13:24 -08:00

shard_map_test.py

Merge pull request #27903 from mattjj:pvary-errors

2025-04-10 09:56:16 -07:00

source_info_test.py

…

source_mapper_test.py

Merge pull request #26369 from justinjfu:sourcemap_win

2025-02-06 12:32:10 -08:00

sourcemap_test.py

…

sparse_bcoo_bcsr_test.py

[JAX] disable flaky parameter permutations for sparse_bcoo_bcsr test.

2025-02-03 16:02:06 -08:00

sparse_test.py

…

sparsify_test.py

[better_errors] Improve error message for lax.switch branches output structure mismatch

2025-01-10 08:03:33 +02:00

stack_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

state_test.py

[Pallas] Fix integer array indexing

2025-04-10 19:10:35 -07:00

stax_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

string_array_test.py

Remove code present to support jaxlib < 0.5.1.

2025-02-26 07:40:40 -05:00

svd_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

transfer_guard_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

tree_util_test.py

Add an option to simplify keystr output and use a custom separator.

2025-01-21 10:18:42 -08:00

typing_test.py

…

unary_ops_accuracy_test.py

Skip jax/tests:unary_ops_accuracy_test when running with older versions of StableHLO.

2025-04-09 13:42:25 -07:00

util_test.py

Add a C++ implementation of a toplogical sort.

2025-03-14 16:04:25 -07:00

version_test.py

Relax jax dependency constraints to be able to install RC wheels

2025-04-07 14:49:55 -07:00

warnings_util_test.py

Switch to a new thread-safe utility for catching warnings.

2025-01-09 11:58:34 -05:00

x64_context_test.py

…

xla_bridge_test.py

Use backend._get_all_devices() to validate devices.

2025-01-24 11:09:16 -08:00

xla_interpreter_test.py

…

xla_metadata_test.py

let XLA metadata be unset in nested dynamic scopes

2025-04-01 20:25:19 -07:00