jax authors b48254e084 Merge pull request #18599 from mattjj:shmap-eager-axis-index ...

PiperOrigin-RevId: 584423708

2023-11-21 14:05:23 -08:00

..

array_api

Add jax.experimental.array_api interface

2023-11-16 14:21:04 -08:00

array_serialization

Add mitigation techniques in the error message when a barrier timeout occurs.

2023-10-27 08:57:39 -07:00

compilation_cache

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

export

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08:00

jax2tf

[jax2tf] Disable test broken by TensorFlow

2023-11-21 01:42:37 -08:00

mosaic

Increase minimum jaxlib version to 0.4.19.

2023-11-17 09:38:31 -08:00

pallas

Avoids some unnecessary dtype casts in Pallas attention to make fwd pass faster. Use num_warps=8 to make bwd faster. Minor fixes to make code more readable.

2023-11-13 15:46:46 -08:00

sparse

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08:00

__init__.py

Copybara import of the project:

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

checkify.py

Document why 'import name as name' is used

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

custom_partitioning.py

Copybara import of the project:

2023-11-06 17:01:02 -08:00

host_callback.py

Cleanup the handling of single- and multi-platform lowering in ModuleContext

2023-10-25 10:40:41 -07:00

jet.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

layout.py

Remove DefaultLayout and make None same as DefaultLayout

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

maps.py

Copybara import of the project:

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

mesh_utils.py

Removing unused description

2023-11-21 12:17:20 -08:00

multihost_utils.py

Switch from flake8 to Ruff

2023-11-15 22:35:52 -05:00

ode.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

pjit.py

Switch from flake8 to Ruff

2023-11-15 22:35:52 -05: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

Stop passing CompileOptions when deserializing.

2023-05-10 16:22:54 -07:00

shard_map.py

[shard_map] add eager axis_index implementation and tests

2023-11-19 11:12:04 -08:00

topologies.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

x64_context.py

MAINT Do not import the config object in JAX internals

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