Yash Katariya 76dec38286 Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env. ...

This would also make it easier to deprecate the `with mesh: pjit` path in the future from user code since the new path would be completely tested.
This will also allow us to remove `resource_env` from JAX and the internal API access of `resource_env.physical_mesh` spread throughout codebases internally and externally.

PiperOrigin-RevId: 735602187

2025-03-10 20:21:02 -07:00

..

_private_mm

Merge pull request #26564 from gspschmid:gschmid/mini_mpmd

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

array_serialization

Remove tensorstore dependency from //jax/experimental/array_serialization:serialization in OSS (see https://github.com/google/tensorstore/issues/218)

2025-02-26 14:47:16 -08:00

colocated_python

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

Cleanup deprecated compilation cache APIs.

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

jax2tf

Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env.

2025-03-10 20:21:02 -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] Test the wgmma_op lowering when a is in registers.

2025-03-10 08:25:43 -07:00

pallas

[Mosaic TPU][Pallas] Add pl.reciprocal

2025-03-07 18:29:30 -08:00

roofline

roofline: Support computing flops for unary ops.

2025-03-06 17:44:36 -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

Under pjit the with mesh: context will use use_mesh(mesh): jit instead of tracking separately using resource_env.

2025-03-10 20:21:02 -07: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

Add public APIs for jax.lax monoidal reductions

2025-02-11 16:00:03 -08: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

Added jax.experimental.multihost_utils.live_devices API.

2025-02-20 10:04:17 -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

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 and other fixes to direct-linearize

2025-03-07 21:02:40 +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