Yash Katariya f95f6a8bdb Rename psum2 to psum_invariant and put it in lax_parallel. We shouldn't expose this to public API and have users use psum instead which will dispatch to psum_invariant when check_rep=True. ...

PiperOrigin-RevId: 745352875

2025-04-08 17:28:59 -07:00

..

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

[Mosaic GPU] Simplify load/store methods now that we have fewer layouts

2025-04-08 07:39:53 -07:00

notebooks

…

pallas

[pallas:mosaic_gpu] Added test for custom pretty-printing rules

2025-04-08 07:59:44 -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

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

2025-04-08 08:33:49 -07: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

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

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

cache_key_test.py

harden cache against jaxlib ver

2025-04-07 23:30:31 +00: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 custom pretty print rule for the unary ops with accuracy s.t. accuracy is not printed if it's None.

2025-04-07 16:25:01 -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

Make changes to shard_map to prepare for setting varying_axes_in_types to True.

2025-04-08 13:47:13 -07: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

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

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

2025-04-08 08:33:49 -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

Relax some test tolerances in N-D FFT tests.

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

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

Merge pull request #27084 from danielsuo:switch-fwd

2025-04-03 01:07:50 -07: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

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

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

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

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

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

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

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

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

2025-03-17 13:41:46 -04: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

Rename psum2 to psum_invariant and put it in lax_parallel. We shouldn't expose this to public API and have users use psum instead which will dispatch to psum_invariant when check_rep=True.

2025-04-08 17:28:59 -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

Disable second order vjp tests in RunStateHypothesisTest.test_vjp if JAX_SKIP_SLOW_TESTS=true

2025-04-08 21:02:55 +00: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 custom pretty print rule for the unary ops with accuracy s.t. accuracy is not printed if it's None.

2025-04-07 16:25:01 -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