jax authors e1cad34522 Add ChunkedCausalMask for Splash Attention to support attention masking similar to Llama4. ...

Llama4 uses (interleaved) chunk attention to support long context.

PiperOrigin-RevId: 746661156

2025-04-11 18:59:07 -07:00

..

_private_mm

Merge pull request #26564 from gspschmid:gschmid/mini_mpmd

2025-02-25 09:15:40 -08:00

array_serialization

Deprecate PositionalSharding and GSPMDSharding

2025-04-11 13:06:43 -07:00

colocated_python

[JAX] Add caching to colocated_python.colocated_cpu_devices()

2025-03-26 15:48:48 -07:00

compilation_cache

Cleanup deprecated compilation cache APIs.

2024-01-12 22:44:48 -08:00

jax2tf

[JAX] Remove calls to jax.dlpack.to_dlpack(), and avoid passing DLPack capsules to jax.dlpack.from_dlpack().

2025-04-11 08:09:41 -07:00

key_reuse

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

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

mosaic

[Mosaic GPU] Implement warp-level thread semantics.

2025-04-10 13:59:22 -07:00

pallas

Add ChunkedCausalMask for Splash Attention to support attention masking similar to Llama4.

2025-04-11 18:59:07 -07:00

roofline

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

slab

Add a variant of safe_map() that has no return value, named foreach().

2025-03-14 07:42:48 -07:00

source_mapper

[Pallas] Plumb compiler flags through source mapper.

2025-03-12 11:19:58 -07:00

sparse

Remove code to support jaxlib < 0.5.3.

2025-03-19 13:40:04 -07:00

__init__.py

add jax.input_saved_vjp to let user pass primal inputs to bwd pass

2025-03-15 00:46:04 +00:00

attrs.py

PR #27576: [attrs] experimental appendattr

2025-03-28 15:21:12 -07:00

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

compute_on.py

Initial commit for jax.experimental.compute_on API.

2024-05-17 15:59:21 -07:00

custom_dce.py

Add an experimental interface for customizing DCE behavior.

2025-01-23 11:38:47 -08:00

custom_partitioning.py

Implement the extension to the custom_partitioning API.

2025-01-08 13:34:47 -08:00

host_callback.py

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

2025-04-08 08:33:49 -07:00

jet.py

Add accuracy field to unary ops

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

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

mesh_utils.py

Add license headers to files that were missing them.

2025-03-27 08:45:15 -07:00

multihost_utils.py

Fix process_allgather of global jax.Arrays with shardy

2025-03-20 09:07:26 -07:00

ode.py

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

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

pjit.py

Remove _pjit_lower_cached cache. We can simplify the caching of jit as we have downstream caches and a cpp cache too.

2024-11-25 11:40:50 -08:00

profiler.py

Add get_profiled_instructions_proto API in jax.experimental.profiler to get profiled instructions proto.

2023-06-29 21:02:05 -07:00

rnn.py

Migrated to mypy 1.14.1 with --allow_redefinition

2025-02-13 15:38:28 +00:00

serialize_executable.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

shard_alike.py

Expose shard_alike via jax.experimental. The API is x, y = shard_like(x, y).

2023-12-19 16:31:33 -08:00

shard_map.py

Rename is_closed to is_open in the shardy shardings

2025-04-11 17:42:34 -07:00

shard.py

Fix the error message to say out_sharding instead of sharding in lax.reshape sharding rule

2025-02-13 07:49:54 -08:00

topologies.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

transfer.py

Remove code present to support jaxlib < 0.5.1.

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

x64_context.py

MAINT Do not import the config object in JAX internals

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

xla_metadata.py

Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.

2024-09-06 16:44:56 -07:00