Yash Katariya 023bc7856b Add registration handler for TPU v5e in mesh_utils. ...

PiperOrigin-RevId: 643092629

2024-06-13 12:52:33 -07:00

..

array_api

Added a link to array-api-tests fix for the tests

2024-05-30 08:38:49 +00:00

array_serialization

Revert: Improve tensorstore I/O efficiency

2024-06-12 15:22:05 -07:00

compilation_cache

Cleanup deprecated compilation cache APIs.

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

export

[export] Rename some API entry points

2024-06-13 06:44:13 +02:00

jax2tf

[export] Rename some API entry points

2024-06-13 06:44:13 +02:00

key_reuse

[key reuse] call key reuse logic directly in dispatch

2024-04-11 17:08:32 -07:00

mosaic

[Mosaic GPU] Don't use enum.StrEnum

2024-06-13 10:14:08 -07:00

pallas

Updated the type annotations of *_spec= parameters of pl.pallas_call

2024-06-11 12:22:00 +01:00

slab

support both static- and dynamic-shaped arguments to djited functions

2024-06-05 22:45:50 -07:00

sparse

[ROCm] Add rocm support in nm.py lowering code

2024-06-05 20:40:35 +00: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

Deprecate XLACompatibleSharding in favor of jax.sharding.Sharding.

2024-06-05 09:07:27 -07:00

host_callback.py

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

jet.py

Add Layout support to jax.jit.

2024-04-05 20:09:34 -07:00

layout.py

Rename layout.AUTO to DeviceLocalLayout.AUTO

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

maps.py

Refactor & test internal deprecation APIs

2024-05-20 17:16:31 -07:00

mesh_utils.py

Add registration handler for TPU v5e in mesh_utils.

2024-06-13 12:52:33 -07:00

multihost_utils.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01: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

Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.

2024-06-03 14:52:50 -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