Dinghua Li 8bf3f47f02 Open source PagedAttention TPU kernel. ...

PiperOrigin-RevId: 620042536

2024-03-28 13:36:02 -07:00

..

array_api

typing: fix incorrect tuple annotations

2024-02-26 10:53:19 -08:00

array_serialization

remove unnecessary if statement

2024-03-20 16:15:15 -07:00

compilation_cache

Cleanup deprecated compilation cache APIs.

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

export

[export] Increase minimum serialization version to 9.

2024-03-27 11:06:23 -07:00

jax2tf

Reverts 0dde8f7f9607d09841ece7125dfc0773c3613fab

2024-03-26 22:26:41 -07:00

key_reuse

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07:00

mosaic

[Mosaic] Remove Python implementation of apply_vector_layout and infer_memref_layout.

2024-01-10 13:00:21 -08:00

pallas

Open source PagedAttention TPU kernel.

2024-03-28 13:36:02 -07:00

sparse

update experimental.sparse to new-style typed keys

2024-03-07 12:40:10 -08:00

__init__.py

Copybara import of the project:

2023-05-24 10:35:37 -07:00

attrs.py

[attrs] fix tracer lifetime bug, fixes #20082

2024-03-05 12:08:44 -08:00

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Deprecate jax.tree_map in favor of jax.tree.map

2024-02-22 11:35:39 -08:00

host_callback.py

Allows for splitting the transpose of a scan into a scan and a map.

2024-03-28 10:54:50 -07:00

jet.py

More numerically stable jax.experimental.jet._integer_pow_taylor.

2024-03-19 09:55:21 -07:00

layout.py

Remove DefaultLayout and make None same as DefaultLayout

2023-11-16 18:01:27 -08:00

maps.py

Register maps module deprecation outside of module

2024-03-19 09:21:48 -07:00

mesh_utils.py

Update the arg description.

2024-03-22 20:10:18 -07:00

multihost_utils.py

Deprecate jax.tree_map in favor of jax.tree.map

2024-02-22 11:35:39 -08:00

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07: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

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08: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 physical_hlo_sharding from TyRules.

2024-03-15 16:02:13 -07: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