George Necula 65450d165e Remove forward compatibility mode for old PRGN custom call on GPU ...

The backend support for the new custom call was added on June 28th.
Also add backwards compatibility test for the new custom call.

PiperOrigin-RevId: 658011228

2024-07-31 08:10:17 -07:00

..

_src

Remove forward compatibility mode for old PRGN custom call on GPU

2024-07-31 08:10:17 -07:00

example_libraries

Run pyupgrade --py310-plus.

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

experimental

[Mosaic GPU] Add a missing return in wait_parity

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

extend

Drop support for mhlo in JAX's public API.

2024-07-30 05:29:52 -07:00

image

…

interpreters

Merge pull request #22263 from hawkinsp:tuples

2024-07-17 09:56:18 -07:00

lax

remove pdot and xeinsum (since xmap is gone)

2024-07-25 21:19:17 +00:00

lib

Switch from flake8 to Ruff

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

nn

PR #21371: [NVIDIA] Add new SDPA API to jax.nn

2024-07-08 06:16:04 -07:00

numpy

[array API] move api metadata into jax.numpy namespace

2024-07-30 12:15:24 -07:00

ops

…

scipy

Copy nn.{softmax,log_softmax} to scipy.special

2024-06-22 09:32:14 -04:00

tools

Added basic support for int2/uint2 dtypes to JAX

2024-07-04 04:13:24 -07:00

__init__.py

Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period).

2024-07-24 10:24:09 -07:00

ad_checkpoint.py

Add memories support to remat.

2023-09-12 20:50:05 -07:00

api_util.py

…

BUILD

[pallas] Use the same primitive run_scoped_p for moth mosaic and mosaic_gpu

2024-07-24 17:14:30 -07:00

cloud_tpu_init.py

…

collect_profile.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

core.py

remove named_shapes (since xmap is now gone)

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

custom_batching.py

…

custom_derivatives.py

custom_vjp symbolic zeros support, take two

2023-04-05 11:17:05 -07:00

custom_transpose.py

…

debug.py

Explicitly re-export all names from jax.debug

2023-04-07 07:24:59 -07:00

distributed.py

…

dlpack.py

…

dtypes.py

Copybara import of the project:

2023-07-24 14:38:20 -07:00

errors.py

[key reuse] define KeyReuseError in jax.errors

2024-03-08 10:59:06 -08:00

export.py

[export] Rename some API entry points

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

flatten_util.py

ravel_pytree now accepted by static linters

2023-03-07 23:35:48 +00:00

monitoring.py

Allow duration event listeners to take extra keyword arguments.

2023-10-04 13:15:49 -07:00

profiler.py

…

py.typed

…

random.py

Finalize deprecation of jax.random.shuffle

2024-07-27 11:21:49 -07:00

sharding.py

Deprecate XLACompatibleSharding in favor of jax.sharding.Sharding.

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

stages.py

rename Specialized to Traced (and specialize to trace)

2024-06-06 17:43:08 -07:00

test_util.py

…

tree_util.py

Export KeyPath and related types to jax.tree_util

2024-07-31 06:41:33 -07:00

tree.py

Add jax.tree module with aliases of jax.tree_util

2024-02-12 13:07:59 -08:00

typing.py

Define jax.typing.DTypeLike

2023-10-10 08:46:36 -07:00

util.py

Add discharge rules for scan with mutable arrays. Move mutable array tests to separate file.

2024-05-02 14:36:16 -04:00

version.py

Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57

2024-07-29 18:44:31 -07:00