Adam Paszke 4c13594bdd [Mosaic GPU] Add a missing return in wait_parity ...

Technically it's not an error to wait twice... but why would you?

PiperOrigin-RevId: 657968887

2024-07-31 05:25:17 -07:00

..

array_api

Merge pull request #22768 from jakevdp:array-api-reshape

2024-07-30 17:15:02 -07:00

array_serialization

Adapt test case

2024-07-15 09:37:59 +00:00

compilation_cache

Cleanup deprecated compilation cache APIs.

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

export

Re-land https://github.com/google/jax/pull/21847

2024-06-13 19:53:53 -07:00

jax2tf

Make the vmap(jit) or vmap(wsc) with a concrete layout error more informative

2024-07-25 18:32:37 -07:00

key_reuse

Run pyupgrade --py310-plus.

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

mosaic

[Mosaic GPU] Add a missing return in wait_parity

2024-07-31 05:25:17 -07:00

pallas

[pallas] More simplification of grid mapping and calling convention

2024-07-29 15:53:47 +02:00

slab

Run pyupgrade --py310-plus.

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

sparse

Re-pack gradients of jax.experimental.sparse.grad() to match original pytrees & test cases

2024-07-29 13:04:05 -04:00

__init__.py

start adding EArray, a jax.Array analog that can contain extended dtypes

2024-04-06 13:09:25 -07:00

attrs.py

fix weak key cache stuff

2024-05-29 17:53:56 +00:00

checkify.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

compute_on.py

Initial commit for jax.experimental.compute_on API.

2024-05-17 15:59:21 -07:00

custom_partitioning.py

Moved the implementation of `custom_partitioning` into jax/_src

2024-07-08 04:31:44 -07:00

host_callback.py

Remove linear parameter from lax.cond_p.

2024-07-01 10:25:42 -04:00

jet.py

Run pyupgrade --py310-plus.

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

layout.py

Rename layout.AUTO to DeviceLocalLayout.AUTO

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

mesh_utils.py

mesh_utils: allow meshes that do not include device at (0, 0, 0).

2024-07-31 01:11:31 -07:00

multihost_utils.py

Roll forward #22076

2024-07-10 12:41:55 -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

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

serialize_executable.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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

remove named_shapes (since xmap is now gone)

2024-07-25 00:54:50 +00:00

topologies.py

[PJRT C API] Supports plugins to register a method to create topology.

2024-02-22 16:55:34 -08:00

x64_context.py

MAINT Do not import the config object in JAX internals

2023-10-18 10:55:13 +01:00