Yash Katariya f8b98993b8 Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings. ...

Also improve dynamic_update_slice sharding error by printing `aval.str_short()` instead of full sharding because it's concise and gives more info than the current error (i.e. it adds shape too to the error message)

Also make some formatting changes in scan lowering to make it easier to debug.

PiperOrigin-RevId: 734542862

2025-03-07 07:01:34 -08:00

..

filecheck

Add public APIs for jax.lax monoidal reductions

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

mosaic

[Mosaic GPU] Add support for tiling stores to refs using small tiling

2025-03-07 05:19:11 -08: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

[Mosaic TPU] Support fp8 upcast to f32

2025-03-06 17:19:15 -08: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

[aot] Add support for as_text(debug_info=True).

2025-01-10 07:59:56 +02:00

api_test.py

Dedup non-ref constants closed in cond branch functions.

2025-03-07 04:01:42 -08:00

api_util_test.py

Consolidate handling of input argument resolution in custom_* APIs.

2024-08-13 00:30:23 -07:00

array_api_skips.txt

array API: improve test coverage

2024-12-05 12:56:07 -08:00

array_api_test.py

Remove deprecated jax.experimental.array_api

2025-01-06 15:19:02 -08:00

array_interoperability_test.py

Increase the minimum NumPy version to v1.25.

2024-12-18 08:18:57 -05:00

array_test.py

Remove _parsed_pspec from everywhere in JAX except for NamedSharding constructor. I'll do that in the next CL since that has a dependency on C++ so needs guards.

2025-02-27 10:17:06 -08:00

attrs_test.py

ref errors

2024-12-18 07:46:14 +00:00

batching_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

blocked_sampler_test.py

[Pallas] Implement block-invariant sampling.

2024-06-24 11:20:39 -07:00

BUILD

[JAX] Disable msan and asan for the profiler test running on nvidia gpu

2025-03-04 10:34:11 -08:00

cache_key_test.py

Remove code that supported jaxlib < 0.5.

2025-01-17 14:22:27 -05:00

checkify_test.py

Don't monkey-patch functions in test_utils to count events for tests.

2024-12-12 09:58:14 -05: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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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 public APIs for jax.lax monoidal reductions

2025-02-11 16:00:03 -08: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] Port the Pallas debug info mechanisms to the new JAX DebugInfo.

2025-02-25 14:43:17 +01:00

debug_nans_test.py

Fix shard_map debug_nan leakage of manual out_avals in the impl rules of jit i.e. impl rule of jit saw a manual out_aval which is not expected. This is a band-aid for now with a TODO to do a proper fix

2025-02-24 10:15:21 -08: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_test.py

add distributed.is_initialized

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

dtypes_test.py

Internal: avoid adding _DimExpr to dtypes._weak_types

2025-01-29 09:11:02 -08: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

Minor bug fixes in error checking

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

errors_test.py

Merge pull request #26279 from MichaelHudgins:tsan-resultstore

2025-02-06 14:55:57 +00:00

experimental_rnn_test.py

Update RNN kernels to use FFI.

2025-02-06 18:27:58 -08:00

export_back_compat_test.py

#sdy Add JAX backwards compatibility test.

2025-03-03 06:01:34 -08: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

#sdy support JAX export tests when Shardy is enabled.

2025-03-03 04:57:06 -08: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

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

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

generated_fun_test.py

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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_to_ir_test.py

[jax2tf] Remove non-native serialization test from jax_to_ir_test

2024-10-07 04:21: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

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

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

lax_control_flow_test.py

Update references to JAX's GitHub repo

2025-03-04 18:31:09 -08: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

jax.numpy reductions: avoid upcast of f16 when dtype is specified by user

2025-02-12 11:49:39 -08:00

lax_numpy_test.py

jax.numpy ndim/shape/size: deprecate non-array input

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

lax_numpy_ufuncs_test.py

Make logaddexp and logaddexp2 into ufuncs

2024-11-18 09:27:36 -08: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

Update jax.scipy.special.gamma and gammasgn to return NaN for negative integer inputs.

2024-11-18 20:33:27 -05: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

Merge pull request #25791 from mattjj:logsumexp-where-grad-nan

2025-01-10 11:27:35 -08:00

lax_test.py

[sharding_in_types] Make the typing checks and sharding rule checks a little bit less strict when the current or aval mesh is empty/unset. Also some more changes as listed below:

2025-02-12 10:03:01 -08: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

Fix linalg.norm to return zero for proper norms of empty matrices.

2025-03-03 15:02:34 -05:00

lobpcg_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

logging_test.py

logging fixes

2024-11-05 13:28:17 -08: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

Remove the skip for test_output_streaming_inside_scan

2025-03-03 14:54:03 -08:00

mesh_utils_test.py

[Take 2] Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh

2024-11-28 09:24:20 -08: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

[better_errors] Cleanup use of DebugInfo.arg_names and result_paths

2025-02-23 08:27:56 +02: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

Improve after review # 2

2025-02-26 04:48:25 +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

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

package_structure_test.py

C++ tree with path API

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

pgle_test.py

Fix test failure in PGLE tests.

2025-02-18 09:44:16 +00:00

pickle_test.py

Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec. We need to do this after sharding-in-types to speed up NamedSharding construction and remove a lot of tech debt and unnecessary complexity.

2025-02-27 08:59:25 -08:00

pjit_test.py

Add a divisibility check so that we make sure that sharding evenly divides the shape (until this restriction is lifted) to make sure we don't create bad shardings.

2025-03-07 07:01:34 -08:00

pmap_test.py

[better_errors] Refactor debug info tests

2025-01-20 20:21:01 +01: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

[xla:cpu] Implement XLA FFI handlers for CPU Jax callbacks.

2025-02-12 13:53:36 -08:00

pytorch_interoperability_test.py

Set PYTHONWARNINGS=error in bazel tests.

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

qdwh_test.py

Bump up tolerance for tests using _testQdwh.

2024-06-14 23:18:20 -07:00

ragged_collective_test.py

Remove code present to support jaxlib < 0.5.1.

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

random_lax_test.py

Add use_high_dynamic_range_gumbel flag which allows sampling gumbel such

2025-03-02 19:42:40 -08:00

random_test.py

Don't monkey-patch functions in test_utils to count events for tests.

2024-12-12 09:58:14 -05:00

roofline_test.py

roofline: Support computing flops for unary ops.

2025-03-06 17:44:36 -08:00

scaled_matmul_stablehlo_test.py

Add apache header.

2025-02-26 15:26:56 -06: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

Make api_test.py work when test cases are run using multiple threads.

2025-01-08 14:09:07 -08: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

PRNGKeyArray.aval should have the correct logical sharding. This required refactoring code so that we don't hit recursion errors.

2025-03-01 18:18:19 -08: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_nm_test.py

[XLA:GPU] Skip small tile sizes for sparse gemms on Ampere as well. Enable the JAX test again that has been failing.

2024-11-11 08:58:15 -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

Avoid 'from jax import config' imports

2024-04-11 13:23:27 -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

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

string_array_test.py

Remove code present to support jaxlib < 0.5.1.

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

svd_test.py

Bump some SVD test tolerances.

2024-07-15 12:54:00 -07:00

transfer_guard_test.py

fix some additional warnings related to #21834

2024-06-13 16:06:14 -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

util_test.py

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

2025-02-09 18:14:33 +02:00

version_test.py

Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.

2025-02-05 10:01:23 -08: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

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