jax authors d42d2e88b4 [Pallas] Interpret dimensions with parallel semantics by traversing the corresponding grid coordinates in randomized order. ...

Note that dynamic grid dimensions with 'parallel' semantics are disallowed. This enables the computation of grid points, with randomized coordinates along 'parallel' dimensions, in Jax/on device.
If randomization of grid dimensions with dynamic sizes (i.e. sizes not known at Jax trace time) were allowed, this would require computing these randomizations on the host/on CPU (where one can have arrays of dynamic shape).

PiperOrigin-RevId: 746365669

2025-04-11 01:54:11 -07:00

..

_src

[Pallas] Interpret dimensions with parallel semantics by traversing the corresponding grid coordinates in randomized order.

2025-04-11 01:54:11 -07:00

example_libraries

…

experimental

[pallas:mgpu] Return types allowed in mgpu.inline_mgpu.

2025-04-10 16:28:34 -07:00

extend

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

image

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

interpreters

Deprecate public export of mlir.custom_call.

2025-04-07 07:58:20 -07:00

lax

Rename pbroadcast to pvary and expose it as jax.lax.pvary.

2025-04-08 16:51:27 -07:00

lib

[JAX] Remove deprecated exports in jax.lib.xla_client.

2025-04-09 14:49:34 -07:00

nn

Enable public doc for scaled dot

2025-03-27 00:05:28 +00:00

numpy

cleanup now that we depend on ml_dtypes>=0.5

2025-03-28 07:44:38 -07:00

ops

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

scipy

add pascal matrix

2025-03-26 02:11:03 +00:00

tools

cleanup now that we depend on ml_dtypes>=0.5

2025-03-28 07:44:38 -07:00

__init__.py

Deprecate jax.tree_* aliases

2025-03-27 10:13:14 -07:00

ad_checkpoint.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

api_util.py

[better_errors] Continue adding debug info to Jaxprs (step 7)

2025-02-09 18:14:33 +02:00

BUILD

Split JaxTestLoader and related classes into a separate file.

2025-04-09 18:45:45 -07:00

cloud_tpu_init.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

collect_profile.py

Remove backward compatibility logic for tool naming.

2025-03-26 23:13:23 -07:00

core.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

custom_batching.py

…

custom_derivatives.py

Reverts 342cb7b99a09180472823a33c7cdad8a8db77875

2025-03-05 10:22:40 -08:00

custom_transpose.py

…

debug.py

…

distributed.py

add distributed.is_initialized

2025-02-18 16:47:19 -08:00

dlpack.py

…

dtypes.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

errors.py

Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class.

2024-09-26 08:39:30 -07:00

export.py

[export] Add support for serialization for some custom PyTree nodes

2024-10-21 11:38:13 +02:00

ffi.py

[xla:python] Add a mechanism for "batch partitioning" of FFI calls.

2025-02-07 09:14:06 -08:00

flatten_util.py

…

monitoring.py

Add scalar event logging function

2025-03-28 23:23:42 +00:00

profiler.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

py.typed

…

random.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

sharding.py

Clean up a number of finalized deprecations

2025-03-26 09:57:19 -07:00

stages.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

test_util.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

tree_util.py

Deprecated jax.tree_util.build_tree

2025-03-21 08:54:38 -07:00

tree.py

Add jax.tree shortcuts for .*_with_path calls, for convenience of users.

2024-12-12 15:13:32 -08:00

typing.py

Replace references to jax.readthedocs.io with docs.jax.dev.

2025-04-08 08:33:49 -07:00

util.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

version.py

Relax jax dependency constraints to be able to install RC wheels

2025-04-07 14:49:55 -07:00