Sergei Lebedev 908ff49e22 [mosaic_gpu] Warmup the kernel when doing CUPTI based profiling ...

Closes #26144.

PiperOrigin-RevId: 730159757

2025-02-23 10:20:37 -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

[JAX] Implement an initial object API for colocated Python

2025-02-21 12:58:25 -08:00

compilation_cache

Cleanup deprecated compilation cache APIs.

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

jax2tf

Reverts 655267609b5589fda8358ae7aaf2eb832036407a

2025-02-21 11:59:44 -08:00

key_reuse

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

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

mosaic

[mosaic_gpu] Warmup the kernel when doing CUPTI based profiling

2025-02-23 10:20:37 -08:00

pallas

[pallas:mosaic_gpu] Added support for binary/comparison ops with WG semantics

2025-02-20 15:06:27 -08:00

roofline

roofline: Handle ClosedJaxpr instances.

2025-02-21 13:19:31 -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

refactor: move lax_numpy indexing routines to their own submodule

2025-02-13 12:03:07 -08: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

If cur_mesh is empty and AxisTypes of Mesh passed to shmap are Explicit, then treat the axes mentioned in auto as explicit too. In other words, "auto" really means "don't convert to manual", ie leave the listed mesh axes as they are, whether explicit or auto type

2025-02-19 21:25:53 -08: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

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