rocm_jax

filecheck

Add public APIs for jax.lax monoidal reductions

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

mosaic

[Mosaic GPU] Define the mosaic_gpu.custom_primitive dialect op.

2025-04-03 00:28:06 -07:00

notebooks

…

pallas

[pallas:mosaic_gpu] Do not specify the default index_map in tests

2025-04-03 22:31:54 -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

Fix custom_transpose when composed with custom_jvp and use_direct_linearize=True.

2025-04-02 15:32:15 -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

jax.numpy: support __jax_array__ in several more functions

2025-04-02 14:28:54 -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

Add auto_axes, explicit_axes and manual_axes properties to Mesh and AbstractMesh

2025-04-03 18:35:28 -07:00

attrs_test.py

PR #27576 : [attrs] experimental appendattr

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

batching_test.py

Change the way that batching.spec_types is updated.

2025-03-15 22:58:44 -07: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

Disabled default env var LOBPCG_EMIT_DEBUG_PLOTS=1

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

cache_key_test.py

[jax:custom_partitioning] Raise an error when Shardy is used but the old sharding propagation callbacks instead of sharding rule are provided.

2025-03-27 13:04:24 -07:00

checkify_test.py

Fix a bug where jit was forwarding inputs to outputs even when donation was True for that inputs. This caused the output to be marked as deleted since the input was being forwarded to the output.

2025-03-26 16:31:11 -07: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

Add accuracy field to unary ops

2025-03-27 17:12:59 -07:00

cudnn_fusion_test.py

Re-enable cudnn_fusion_test on A100.

2024-11-01 15:48:07 +00:00

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

[better_errors] Fix the handling of kwargs for debug_info.

2025-04-02 10:32:38 +01: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

Remove code present to support jaxlib < 0.5.1.

2025-02-26 07:40:40 -05: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

Merge pull request #26279 from MichaelHudgins:tsan-resultstore

2025-02-06 14:55:57 +00: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

Remove legacy GPU kernels for QR decomposition.

2025-04-02 08:10:11 -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

[export] Add support for override_lowering_rules to jax.export.

2025-04-03 16:13:16 +01: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

Remove code present to support jaxlib < 0.5.1.

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

fft_test.py

Relax some test tolerances in N-D FFT tests.

2024-12-20 03:21:53 -08:00

for_loop_test.py

Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.

2024-10-29 11:04:31 -07:00

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

Add OOB checks to jax.numpy array indexing

2025-04-01 19:17:57 -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

Adding support for copy_p primitive to jet.

2024-11-07 17:06:02 -08:00

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

Merge pull request #27084 from danielsuo:switch-fwd

2025-04-03 01:07:50 -07:00

lax_metal_test.py

adopt jax.extend.backend

2024-10-07 12:27:35 -07:00

lax_numpy_einsum_test.py

…

lax_numpy_indexing_test.py

Improved errors when indexing with floats

2025-02-28 15:04:07 -08:00

lax_numpy_operators_test.py

Implement jax.numpy.spacing

2024-10-03 10:40:39 -07:00

lax_numpy_reducers_test.py

Skip failing tests on TPU v6+

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

lax_numpy_test.py

jnp.array no longer accepts None

2025-04-02 14:58:51 -07: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

Expose jax.lax.shape_as_value().

2025-03-25 09:12:31 -07:00

lax_vmap_op_test.py

…

lax_vmap_test.py

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

2024-12-17 10:12:34 -08:00

layout_test.py

Make sure default layout is None for input and output layout in all codepaths

2025-02-27 14:26:25 -08: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

add pascal matrix

2025-03-26 02:11:03 +00: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

Enable test_scan_offload in memories_test.

2025-04-01 14:26:26 -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

…

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] add vmap rule for mutable_array_p, very basic test

2025-04-01 20:18:32 +00: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

Add identity activation

2025-03-23 15:13:17 -07: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

C++ tree with path API

2024-11-30 21:26:48 -08:00

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 an out_sharding option to jax.random.permutation

2025-04-03 16:21:45 -07:00

pmap_test.py

Add accuracy field to unary ops

2025-03-27 17:12:59 -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

[jaxlib] Fix asan tests for subbyte types in CPU/GPU callbacks.

2025-04-02 03:40:20 -07:00

pytorch_interoperability_test.py

Set PYTHONWARNINGS=error in bazel tests.

2024-09-24 12:30:11 -07:00

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

Add replace option to random.categorical to enable sampling without replacement.

2025-03-17 13:41:46 -04:00

random_test.py

implement nbytes for PRNGKeyArray

2025-03-28 10:54:48 -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

Fix dtype bug in jax.scipy.fft.idct

2024-09-25 12:55:43 -04:00

scipy_interpolate_test.py

…

scipy_ndimage_test.py

…

scipy_optimize_test.py

Fix some warnings causing CI failures on ARM.

2024-09-24 17:25:26 -07:00

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

Bump up tolerance in ShardMapSystematicTest.test_vmap_closure for GPUs.

2025-04-03 08:44:15 -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

[better_errors] Continue adding debug info to Jaxprs (step 7)

2025-02-09 18:14:33 +02: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

Add accuracy field to unary ops

2025-03-27 17:12:59 -07:00

util_test.py

Add a C++ implementation of a toplogical sort.

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

version_test.py

Modify version test to consider "rc" versions as well

2025-03-10 11:10:18 -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