Yash Katariya 0a72e856cf Add **experimental** with_dll_constraint API. This is for cases when the users wants to let SPMD decide the sharding. ...

But this is a contradiction since layouts apply to device local shape and without knowing the sharding, you can't decide the layout. But there are cases where you don't care what the sharding is, you just want to force a row-major layout (for example). **This API should only be used for those cases**.

PiperOrigin-RevId: 744888557

2025-04-07 16:21:58 -07:00

..

_private_mm

Merge pull request #26564 from gspschmid:gschmid/mini_mpmd

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

array_serialization

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -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

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

key_reuse

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

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

mosaic

[Mosaic GPU] Add scaffolding for a new lowering "axis" (UserThreadSemantics), in addition to the existing ThreadSemantics (renamed to LoweringSemantics).

2025-04-07 14:47:59 -07:00

pallas

[Mosaic GPU] Add scaffolding for a new lowering "axis" (UserThreadSemantics), in addition to the existing ThreadSemantics (renamed to LoweringSemantics).

2025-04-07 14:47:59 -07:00

roofline

Support conv unfused_flops in roofline.

2025-04-02 14:08:30 -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

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

2024-10-08 04:22:20 -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

[shard-map] in eager shmap, handle all rep rule output cases

2025-04-07 19:31:07 +00: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