Dan Foreman-Mackey f93c2a1aa5 Add and test support for partitioning of batch dimensions in lax.linalg. ...

On CPU and GPU, almost all of the primitives in lax.linalg are backed by custom calls that support simple semantics when batch dimensions are sharded. Before this change, all linalg operations on CPU and GPU will insert an `all-gather` before being executed when called on sharded inputs, even when that shouldn't be necessary. This change adds support for this type of partitioning, to cover a wide range of use cases.

There are a few remaining GPU ops that don't support partitioning either because they are backed by HLO ops that don't partition properly (Cholesky factorization and triangular solves), or because they're still using descriptors with problem dimensions in kernel. I'm going to fix these in follow up changes.

PiperOrigin-RevId: 731732301

2025-02-27 08:16:16 -08:00

..

filecheck

Add public APIs for jax.lax monoidal reductions

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

mosaic

Remove code present to support jaxlib < 0.5.1.

2025-02-26 07:40:40 -05: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:triton] Generate more efficient code for loading contiguous slices of int4 values.

2025-02-27 01:57:47 -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

Add jax.copy_to_host_async(tree).

2025-02-27 01:22:15 -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

Rolling back a commit that caused a 50-90% performance regression in most MaxText workloads.

2025-02-26 16:57:18 -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

Add and test support for partitioning of batch dimensions in lax.linalg.

2025-02-27 08:16:16 -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

Remove code present to support jaxlib < 0.5.1.

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

Mark error_check_test as thread-unsafe.

2025-02-24 07:36:18 -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

Some lax.linalg housekeeping.

2025-02-10 08:27:18 -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

skip tests with extra requirements

2025-02-05 01:48:28 +00: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

Fix cudnn version skipping in fused_attention_stablehlo_test

2025-02-27 02:19:43 -08: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

Fix failures in TSAN free threading CI.

2025-02-26 06:04:26 -05: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

Indexing: avoid dynamic_slice when mode='clip'

2025-01-14 11:20:50 -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

Reverts f3fade3b70443b6cf87f01f360e6a1cb85d4b1fb

2025-02-27 03:26:37 -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

Remove code present to support jaxlib < 0.5.1.

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

lazy_loader_test.py

…

linalg_sharding_test.py

Add and test support for partitioning of batch dimensions in lax.linalg.

2025-02-27 08:16:16 -08:00

linalg_test.py

Remove code present to support jaxlib < 0.5.1.

2025-02-26 07:40:40 -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 code present to support jaxlib < 0.5.1.

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

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

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

nn_test.py

Merge pull request #25766 from carlosgmartin:nn_initializers_variance_scaling_mode_fan_geo_avg

2025-01-31 15:41:50 -08: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

Avoid 'from jax import config' imports

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

pjit_test.py

Fix a test failure under multi-threading.

2025-02-27 06:07:05 -05: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

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

2025-01-10 06:58: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

Extend random.orthogonal to semi-orthogonal matrices. Simplify initializers.orthogonal by using it.

2025-02-26 16:39:45 -05: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

Add unfused_hbm usage to binary ops and dot_general.

2025-02-25 16:10:25 -08: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

#sdy close any partially sharded dimensions if using auto axes in a shard_map.

2025-02-27 07:53:18 -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