Adam Paszke 21598d02e5 [Mosaic GPU] Add support for non-multicast .cta_group::2 async_copies ...

This instruction is particularly useful for collective MMA, since it lets us
easily report on the progress of async copies from both blocks in the single
block that will be performing the MMA.

PiperOrigin-RevId: 725618793

2025-02-11 07:13:35 -08:00

..

array_serialization

[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

colocated_python

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

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

compilation_cache

Cleanup deprecated compilation cache APIs.

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

jax2tf

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

2025-02-11 11:28:58 +01:00

key_reuse

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

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

mosaic

[Mosaic GPU] Add support for non-multicast .cta_group::2 async_copies

2025-02-11 07:13:35 -08:00

pallas

[Pallas] Add pipeline mode to pltpu

2025-02-10 02:36:44 -08:00

roofline

Cleanup: toward merging core.concrete_aval & xla.abstractify

2024-12-17 09:27:00 -08:00

slab

Remove ConcreteArray from JAX. It's easy to do trace-time concretization without it.

2024-10-31 14:06:54 -07:00

source_mapper

Propagate source ranges in location.

2025-02-03 17:32:59 -08:00

sparse

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

2025-02-11 11:28:58 +01:00

__init__.py

add jax.experimental.primal_tangent_dtype helper

2024-09-21 20:35:20 +00:00

attrs.py

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

2025-02-11 11:28:58 +01: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

Remove remaining implementations of jax.experimental.host_callback.call.

2024-10-08 04:22:20 -07:00

jet.py

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

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

layout.py

Rename layout.AUTO to DeviceLocalLayout.AUTO

2024-04-03 17:23:35 -07:00

mesh_utils.py

Create jax.make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str], devices: Sequence[jax.Device] | None = None) API to make it easier to create a mesh and reduce a ton of boilerplate.

2024-09-03 14:32:03 -07:00

multihost_utils.py

Match the behavior on single host wrt multi-host if tiled=False. Fixes https://github.com/jax-ml/jax/issues/25783

2025-01-08 13:31:35 -08: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

[ROCm] Implement RNN support

2025-01-14 19:04:49 -06: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

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

2025-02-11 11:28:58 +01:00

shard.py

Replace Hidden/Visible/Collective AxisTypes names with Auto/Explicit/Manual.

2025-01-24 23:21:13 -08:00

topologies.py

Cleanup: fix unused imports & mark exported names

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

transfer.py

Move transfer python bindings into jax.

2025-01-23 17:52:57 -08: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