Adam Paszke bb2e2303d7 [Pallas:MGPU] Treat each warpgroup as a single logical thread. ...

As an extra minor change, we now disallow specifying the predicate when uniform is
unset, as that implies that we're going to use two different mechanisms to select
a single thread.

PiperOrigin-RevId: 689289365

2024-10-24 01:54:10 -07:00

..

array_api

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

array_serialization

Cleanup: fix unused imports & mark exported names

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

compilation_cache

Cleanup deprecated compilation cache APIs.

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

export

[export] Add support for serialization for some custom PyTree nodes

2024-10-21 11:38:13 +02:00

jax2tf

Add test for jax2tf conversion of dot general with algorithm.

2024-10-23 08:34:52 -07:00

key_reuse

Cleanup: fix unused imports & mark exported names

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

mosaic

[Pallas:MGPU] Treat each warpgroup as a single logical thread.

2024-10-24 01:54:10 -07:00

pallas

[Pallas:MGPU] Fix the implementation of WGMMA with transposed RHS

2024-10-18 10:31:42 -07:00

slab

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

sparse

[sharding_in_types][Take 2] Add out_type argument to einsum and dot_general to allow specifying for the output type. Right now, it only accept a NamedSharding but in the future we can allow a polymorphic type of: jax.ShapeDtypeStruct | Sharding | Layout.

2024-10-22 13:10:43 -07:00

__init__.py

add jax.experimental.primal_tangent_dtype helper

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

attrs.py

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -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_partitioning.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

host_callback.py

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

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

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

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

Cleanup: fix unused imports & mark exported names

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

ode.py

Update jnp.clip to Array API 2023 standard

2024-04-04 22:55:10 +00:00

pjit.py

Remove a number of deprecated APIs

2023-12-18 10:08:47 -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

Simplify and consolidate dot algorithm control in lax.

2024-10-07 13:21:34 -07: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

Cleanup: fix unused imports & mark exported names

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

topologies.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07: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