rocm_jax

filecheck

Add public APIs for jax.lax monoidal reductions

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

mosaic

[Mosaic GPU] Add support for stores to TMEM

2025-03-28 05:12:08 -07:00

notebooks

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

pallas

[pallas] Removed pl.device_id. Use lax.axis_index instead.

2025-04-02 15:45:52 -07:00

testdata

…

third_party/scipy

Fix CI test failure in line_search_test.

2024-05-15 07:38:20 -07:00

ann_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

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

jnp.concat and friends: support __jax_array__

2025-04-02 13:17:59 -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

Remove some long-stale version guards.

2025-03-21 13:23:26 -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

Bump minimum jaxlib version to v0.4.30.

2024-06-18 12:35:08 -04:00

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

Set the value of config.jax_platforms directly instead of setting the env variable JAX_PLATFORMS. Setting JAX_PLATFORMS doesn’t do anything because config.jax_platforms is already initialized at that point.

2024-07-22 15:09:15 -07:00

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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

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

Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57

2024-07-29 18:44:31 -07:00

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

Lower lax.ragged_dot_general to chlo.ragged_dot in some cases on tpu.

2025-03-21 17:36:32 -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

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

Add CUDA custom call example as a JAX test

2024-06-17 15:21:49 -04:00

heap_profiler_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

image_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

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

add pjit forwarding rule

2024-05-25 17:46:01 +00:00

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

fix a linearize-of-remat-of-while_loop-fixpoint bug

2025-03-23 03:50:55 +00:00

lax_metal_test.py

adopt jax.extend.backend

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

lax_numpy_einsum_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

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

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

lax_scipy_sparse_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

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

Disable test_source_file_prefix_removal test because there is cross-contamination of metadata information from different call sites because of cached jitted functions

2024-07-12 12:07:24 -07:00

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

Allow duration event listeners to take extra keyword arguments.

2023-10-04 13:15:49 -07:00

mosaic_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

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

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

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.bits

2025-04-02 17:19:57 -07:00

pmap_test.py

Add accuracy field to unary ops

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

polynomial_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

pretty_printer_test.py

Compute source maps when pretty-printing jaxprs.

2024-05-06 15:45:25 -04:00

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

Merge pull request #26944 from wenscarl:wenscarl/nvfp4

2025-03-13 13:30:46 -07: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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

scipy_ndimage_test.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

scipy_optimize_test.py

Fix some warnings causing CI failures on ARM.

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

scipy_signal_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

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

Add vma rules for pmin and pmax

2025-03-28 16:55:16 -07:00

source_info_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -07:00

source_mapper_test.py

Merge pull request #26369 from justinjfu:sourcemap_win

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

sourcemap_test.py

Merge pull request #21096 from gspschmid:gschmid/sourcemaps

2024-05-08 05:44:08 -07:00

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

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

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

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

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

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

xla_bridge_test.py

Use backend._get_all_devices() to validate devices.

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

xla_interpreter_test.py

Move some utilities out of dispatch.py next to their users, add more types.

2023-08-08 10:52:11 -07:00

xla_metadata_test.py

let XLA metadata be unset in nested dynamic scopes

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