Sergey Kozub aebe82a78f Add JAX API that provides sparse matmul support (2:4 structured sparsity) ...

Usage:
from jax.experimental.sparse import nm
res = nm.nm_spmm(lhs, rhs, nm.nm_pack(mask))

where:
lhs.shape = [M, K/2]
rhs.shape = [K, N]
`mask` has the same shape as `lhs` with boolean type

If batch dimensions are present, the `dimension_numbers` argument has to be set to:
((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, rhs_batch_dims))

The lowering only works on nVidia GPUs, that provide hardware support for sparse dots.

PiperOrigin-RevId: 627640553

2024-04-24 01:06:19 -07:00

..

filecheck

tests: access tree utilities via jax.tree.*

2024-02-26 14:17:18 -08:00

mosaic

Initial commit for Mosaic GPU

2024-04-18 04:04:10 -07:00

notebooks

Deprecate arr.device_buffer and arr.device_buffers

2023-12-06 10:20:29 -08:00

pallas

[Pallas TPU] Generalize while_loop lowering in Pallas -> Mosaic.

2024-04-18 11:03:52 -07:00

testdata

…

third_party/scipy

Avoid 'from jax import config' imports

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

ann_test.py

Avoid 'from jax import config' imports

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

aot_test.py

Avoid 'from jax import config' imports

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

api_test.py

Fix donation with kwargs. The problem is that pytrees sort dictionaries by default. So if we create the donation vector with original kwargs order, it won't match the aval order (which is created by sorting kwargs i.e. dict) and we end up donating the wrong input.

2024-04-23 14:50:04 -07:00

api_util_test.py

Replace donation_vector's logic with donation_vector_with_in_tree which is now deleted

2024-04-23 17:38:30 -07:00

array_api_test.py

Add new cumulative_sum function to numpy and array_api

2024-04-16 19:57:55 +00:00

array_interoperability_test.py

[PJRT C API] Add stream extension to support DLPack and implement this extension in CUDA plugin.

2024-04-19 10:41:55 -07:00

array_test.py

If callback returns a fully replicated global array, return it as is.

2024-04-14 14:35:57 -07:00

attrs_test.py

[attrs] fix tracer lifetime bug, fixes #20082

2024-03-05 12:08:44 -08:00

batching_test.py

Avoid 'from jax import config' imports

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

BUILD

Add JAX API that provides sparse matmul support (2:4 structured sparsity)

2024-04-24 01:06:19 -07:00

cache_key_test.py

Strip device_assignment on GPU platform.

2024-03-19 01:50:39 -07:00

checkify_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

clear_backends_test.py

Avoid 'from jax import config' imports

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

compilation_cache_test.py

Adds meaningful function names for better debugging.

2024-04-23 09:45:57 -07:00

core_test.py

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07:00

custom_linear_solve_test.py

Avoid 'from jax import config' imports

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

custom_object_test.py

Avoid 'from jax import config' imports

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

custom_root_test.py

Avoid 'from jax import config' imports

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

debug_nans_test.py

Avoid 'from jax import config' imports

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

debugger_test.py

jax.debug.callback now passes arguments as jax.Arrays

2024-04-19 13:57:18 -07:00

debugging_primitives_test.py

jax.debug.callback now passes arguments as jax.Arrays

2024-04-19 13:57:18 -07:00

deprecation_test.py

Fix tests that mistakenly used assertRaises(..., msg=...) to match a message.

2023-10-06 08:07:49 -04:00

dtypes_test.py

Added int4 and uint4 to dtype-specific tests

2024-04-18 15:20:20 +01:00

dynamic_api_test.py

Avoid 'from jax import config' imports

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

errors_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

experimental_rnn_test.py

MAINT Do not import the config object in JAX internals

2023-10-18 10:55:13 +01:00

export_back_compat_test.py

Bumped the minimum jaxlib to 0.4.23

2024-04-13 08:18:33 +01:00

export_harnesses_multi_platform_test.py

[ROCm]: Add rocm as platform in export_harnessess_multi_platform_test

2024-03-20 16:24:06 +00:00

export_test.py

Add Layout support to jax.jit.

2024-04-05 20:09:34 -07:00

extend_test.py

Avoid 'from jax import config' imports

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

fft_test.py

Updated

2024-02-05 18:01:48 -05:00

for_loop_test.py

Avoid 'from jax import config' imports

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

fused_attention_stablehlo_test.py

Support BNTH input formats

2024-04-03 20:48:37 +00:00

generated_fun_test.py

Avoid 'from jax import config' imports

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

gfile_cache_test.py

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

gpu_memory_flags_test.py

Exclude test_gpu_memory_allocation from pytest execution.

2023-12-20 14:53:33 -08:00

heap_profiler_test.py

Avoid 'from jax import config' imports

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

host_callback_test.py

Avoid 'from jax import config' imports

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

host_callback_to_tf_test.py

[callback] Add a flag to implement host_callback in terms of io_callback.

2024-04-05 08:51:30 +01:00

image_test.py

Avoid 'from jax import config' imports

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

infeed_test.py

Avoid 'from jax import config' imports

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

jax_jit_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

jax_to_ir_test.py

Added int4 and uint4 to dtype-specific tests

2024-04-18 15:20:20 +01:00

jaxpr_effects_test.py

Remove internal references to deprecated jax.experimental.maps

2024-03-19 09:24:52 -07:00

jaxpr_util_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

jet_test.py

Avoid 'from jax import config' imports

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

key_reuse_test.py

[key reuse] handle reuse of closed-over constants

2024-04-11 15:39:45 -07:00

lax_autodiff_test.py

Avoid 'from jax import config' imports

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

lax_control_flow_test.py

Let xs=None by default in lax.scan.

2024-04-23 17:26:23 -04:00

lax_metal_test.py

CI: run the metal CI nightly & on PRs touching the config file

2024-04-11 11:18:27 -07:00

lax_numpy_einsum_test.py

Avoid 'from jax import config' imports

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

lax_numpy_indexing_test.py

jax.numpy: implement scalar boolean indexing

2024-02-09 11:00:00 -08:00

lax_numpy_operators_test.py

Fixed hypot bug on nan/inf pairings, began deprecation of non-real values

2024-04-15 17:56:16 +00:00

lax_numpy_reducers_test.py

Add support for copy kwarg in astype to match Array API

2024-04-22 16:25:37 +00:00

lax_numpy_test.py

Merge pull request #20862 from jakevdp:select_n

2024-04-23 18:53:05 -07:00

lax_numpy_ufuncs_test.py

Avoid 'from jax import config' imports

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

lax_numpy_vectorize_test.py

Avoid 'from jax import config' imports

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

lax_scipy_sparse_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

lax_scipy_special_functions_test.py

Add jax.scipy.special.gammasgn

2024-04-18 16:14:55 -07: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

Avoid 'from jax import config' imports

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

lax_test.py

Workaround mpmath 1.3 bugs in tan and tanh evaluation at infinities

2024-04-10 18:26:07 +03:00

lax_vmap_op_test.py

Avoid 'from jax import config' imports

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

lax_vmap_test.py

Avoid 'from jax import config' imports

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

layout_test.py

Add layout support to make_array_from_callback.

2024-04-15 12:38:34 -07:00

lazy_loader_test.py

…

linalg_test.py

Bumped the minimum jaxlib to 0.4.23

2024-04-13 08:18:33 +01:00

lobpcg_test.py

Avoid 'from jax import config' imports

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

logging_test.py

Avoid 'from jax import config' imports

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

memories_test.py

[XLA:TPU] Support scan loops for parameter input and output streaming in host offloading.

2024-04-21 20:03:11 -07:00

mesh_utils_test.py

Adding expectations to e2e test with expectations for mesh creation allowing splitting physical axes.

2024-03-21 18:58:06 -07:00

metadata_test.py

Avoid 'from jax import config' imports

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

mock_gpu_test.py

Avoid 'from jax import config' imports

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

Avoid 'from jax import config' imports

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

multibackend_test.py

Avoid 'from jax import config' imports

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

multiprocess_gpu_test.py

Avoid 'from jax import config' imports

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

name_stack_test.py

Avoid 'from jax import config' imports

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

nn_test.py

softmax: deprecate initial argument & always set to -inf internally

2024-04-10 10:23:21 -07:00

ode_test.py

Avoid 'from jax import config' imports

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

optimizers_test.py

Avoid 'from jax import config' imports

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

package_structure_test.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

pgle_test.py

Avoid 'from jax import config' imports

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

pickle_test.py

Avoid 'from jax import config' imports

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

pjit_test.py

Skip test_spmd_preserves_input_sharding_vmap_grad unless xla_extension_version >= 258

2024-04-18 09:37:57 -07:00

pmap_test.py

Guard test that requires two devices with device_count() check.

2024-04-01 12:32:54 -07:00

polynomial_test.py

Avoid 'from jax import config' imports

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

profiler_test.py

Avoid 'from jax import config' imports

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

python_callback_test.py

relax a side-effects test that was erroneously checking against a canonical order

2024-04-15 17:13:35 -07:00

pytorch_interoperability_test.py

Enable some more C API tests.

2024-04-22 09:38:59 -07:00

qdwh_test.py

MAINT Migrate remaining internal/test modules to use state objects

2023-10-12 17:32:15 +01:00

random_lax_test.py

remove test of deprecated jax.random.shuffle API

2024-04-10 08:19:12 -07:00

random_test.py

Finalize the deprecation of the arr.device() method

2024-04-08 19:04:15 -07:00

scipy_fft_test.py

Avoid 'from jax import config' imports

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

scipy_interpolate_test.py

Avoid 'from jax import config' imports

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

scipy_ndimage_test.py

Avoid 'from jax import config' imports

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

scipy_optimize_test.py

Avoid 'from jax import config' imports

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

scipy_signal_test.py

Avoid 'from jax import config' imports

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

scipy_spatial_test.py

Avoid 'from jax import config' imports

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

scipy_stats_test.py

Merge pull request #20885 from rajasekharporeddy:test_branch4

2024-04-23 13:29:40 -07:00

shape_poly_test.py

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07:00

shard_alike_test.py

Avoid 'from jax import config' imports

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

shard_map_test.py

Support auto in shard_map.

2024-04-22 14:29:35 -07:00

source_info_test.py

Avoid 'from jax import config' imports

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

sparse_bcoo_bcsr_test.py

Avoid 'from jax import config' imports

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

sparse_nm_test.py

Add JAX API that provides sparse matmul support (2:4 structured sparsity)

2024-04-24 01:06:19 -07:00

sparse_test.py

Avoid 'from jax import config' imports

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

sparsify_test.py

Avoid 'from jax import config' imports

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

stack_test.py

Avoid 'from jax import config' imports

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

state_test.py

Add a zeros rule for mutable arrays and test it using a custom vjp.

2024-04-12 15:22:07 -04:00

stax_test.py

Avoid 'from jax import config' imports

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

svd_test.py

Fix bug in rank-deficient fix-up code: Do not zero out the corresponding column of u_out if a diagonal entry of r is exactly zero.

2024-04-18 09:20:48 -07:00

transfer_guard_test.py

Avoid 'from jax import config' imports

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

tree_util_test.py

Add builtin cc dataclass pytree node for performance.

2024-04-23 14:14:49 -07:00

typing_test.py

Upgrade remaining sources to Python 3.9

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

util_test.py

Avoid 'from jax import config' imports

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

version_test.py

Add more extensive tests for version strings

2023-09-08 11:33:49 -07:00

x64_context_test.py

Avoid 'from jax import config' imports

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

xla_bridge_test.py

Change xla_bridge_test to expect a bytes FDO profile instead of a string in future jaxlib versions.

2024-03-06 13:50:36 -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

xmap_test.py

Avoid 'from jax import config' imports

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