rocm_jax

clusters

Return False in is_env_present if importing kubernetes leads to a ModuleNotFoundError

2025-04-16 15:15:27 -07:00

cudnn

Merge pull request #28018 from Cjkkkk:disable_packed_layout_at_ampere

2025-04-16 10:54:25 -07:00

debugger

Cleanup: fix unused imports & mark exported names

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

export

Deprecate PositionalSharding and GSPMDSharding

2025-04-11 13:06:43 -07:00

extend

Move jex.ffi to jax.ffi.

2024-12-29 13:06:19 +00:00

image

Migrated to mypy 1.14.1 with --allow_redefinition

2025-02-13 15:38:28 +00:00

internal_test_util

Remove legacy CPU custom call kernels that have been unused since v0.4.34.

2025-04-11 03:17:19 -07:00

interpreters

Remove jax_varying_axes_in_types config and rewrite from shard_map_p

2025-04-16 22:27:50 -07:00

lax

Remove jax_varying_axes_in_types config and rewrite from shard_map_p

2025-04-16 22:27:50 -07:00

lib

Split weakref_lru_cache into its own extension.

2025-04-08 13:34:47 -07:00

nn

Clean up softmax initial deprecation

2025-04-15 14:36:56 -07:00

numpy

jnp.power: better docs for invalid input

2025-04-14 10:42:29 -07:00

ops

Minor improvements to doc for jax.nn.logsumexp.

2025-04-13 15:17:11 -04:00

pallas

[Pallas][Mosaic][TPU] Add disable_bounds_checks compiler params

2025-04-16 01:01:27 -07:00

scipy

add pascal matrix

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

state

[Pallas] Fix integer array indexing

2025-04-10 19:10:35 -07:00

third_party

refactor: move scalar type defs out of lax_numpy.py

2025-02-06 14:48:10 -08:00

__init__.py

…

abstract_arrays.py

[sharding_in_types] Make the typing checks and sharding rule checks a little bit less strict when the current or aval mesh is empty/unset. Also some more changes as listed below:

2025-02-12 10:03:01 -08:00

ad_checkpoint.py

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

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

ad_util.py

[NFC] Rename standard_insert_pbroadcast to standard_insert_pvary

2025-04-15 11:02:45 -07:00

api_util.py

Fix bug in argnums_partial_except when static_argnums is unsorted.

2025-04-16 16:18:10 -04:00

api.py

jax.jit: deprecate non-standard call signature.

2025-04-14 10:13:05 -07:00

array.py

Moved the jax.Array baseclass to C++

2025-04-03 08:28:02 -07:00

basearray.py

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

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

basearray.pyi

Moved the jax.Array baseclass to C++

2025-04-03 08:28:02 -07:00

blocked_sampler.py

Make blocked_fold_in consistent when the block sizes induce padding

2025-03-18 09:12:11 -07:00

cache_key.py

harden cache against jaxlib ver

2025-04-07 23:30:31 +00:00

callback.py

Finalize deprecation of vectorized argument in callbacks.

2025-04-14 07:43:59 -07:00

checkify.py

Make changes to shard_map to prepare for setting varying_axes_in_types to True.

2025-04-08 13:47:13 -07:00

cloud_tpu_init.py

Warn the user if transparent huge pages aren't enabled.

2025-03-10 10:37:58 -07:00

compilation_cache_interface.py

Fix detection of epath

2025-03-04 11:44:27 -05:00

compilation_cache.py

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

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

compiler.py

AutoPGLE: force-disable graphs less

2025-03-31 18:01:56 +00:00

compute_on.py

Minor format cleanup

2025-02-15 04:56:27 -08:00

config.py

Remove jax_varying_axes_in_types config and rewrite from shard_map_p

2025-04-16 22:27:50 -07:00

core.py

Remove jax_varying_axes_in_types config and rewrite from shard_map_p

2025-04-16 22:27:50 -07:00

custom_api_util.py

…

custom_batching.py

Improve error messages when input argument resolution fails in custom_* APIs.

2025-03-04 10:31:35 -05:00

custom_dce.py

Improve error messages when input argument resolution fails in custom_* APIs.

2025-03-04 10:31:35 -05:00

custom_derivatives.py

Add JVP rule for linear_call.

2025-04-10 09:12:01 -04:00

custom_partitioning_sharding_rule.py

[jax:custom_partitioning] Support SdyShardingRule with multiple leading

2025-01-23 08:20:38 -08:00

custom_partitioning.py

[jax:custom_partitioning] Raise an error when Shardy is used but the old sharding propagation callbacks instead of sharding rule are provided.

2025-03-27 13:04:24 -07:00

custom_transpose.py

Fix custom_transpose when composed with custom_jvp and use_direct_linearize=True.

2025-04-02 15:32:15 -04:00

debugging.py

Rename is_closed to is_open in the shardy shardings

2025-04-11 17:42:34 -07:00

deprecations.py

jax.jit: deprecate non-standard call signature.

2025-04-14 10:13:05 -07:00

dispatch.py

Deprecate PositionalSharding and GSPMDSharding

2025-04-11 13:06:43 -07:00

distributed.py

Don't set memory_kind to None if the mesh is AbstractMesh and the

2025-04-04 14:51:12 -07:00

dlpack.py

Pulled DLDeviceType to XLA backend mapping into a global

2025-01-29 11:38:50 -08:00

dtypes.py

cleanup now that we depend on ml_dtypes>=0.5

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

earray.py

PRNGKeyArray.aval should have the correct logical sharding. This required refactoring code so that we don't hit recursion errors.

2025-03-01 18:18:19 -08:00

effects.py

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

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

environment_info.py

Update device info printed in print_environment_info to print the TPU kind also along with global and local device count.

2024-10-04 09:20:27 -07:00

error_check.py

Minor docstring updates for AOT wrappers in error checking

2025-03-31 15:55:18 -07:00

errors.py

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

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

ffi.py

[NFC] Rename standard_insert_pbroadcast to standard_insert_pvary

2025-04-15 11:02:45 -07:00

flatten_util.py

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

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

hardware_utils.py

Warn the user if transparent huge pages aren't enabled.

2025-03-10 10:37:58 -07:00

jaxpr_util.py

[better_errors] Expand the tests for debug_info

2025-01-22 16:49:16 +01:00

lax_reference.py

Add int4/uint4 support to bitcast_convert_type

2024-12-20 12:45:24 -08:00

layout.py

Add **experimental** with_dll_constraint API. This is for cases when the users wants to let SPMD decide the sharding.

2025-04-07 16:21:58 -07:00

lazy_loader.py

lazy_loader.attach now only imports the submodule once on first access

2024-08-28 15:37:28 -07:00

linear_util.py

[explain_cache_miss] Add to explanations the duration of the missed function call

2025-04-14 16:08:24 +03:00

logging_config.py

Adding JAX_LOGGING_LEVEL configuration option

2024-11-05 09:56:46 -08:00

lru_cache.py

Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1

2025-02-14 12:01:55 -08:00

mesh_utils.py

Add license headers to files that were missing them.

2025-03-27 08:45:15 -07:00

mesh.py

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

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

monitoring.py

Add scalar event logging function

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

named_sharding.py

Remove _manual_axes from NamedSharding since we can now track the manual axes on the mesh.

2025-04-16 21:49:53 -07:00

op_shardings.py

Use "x is y" rather than "id(x) == id(y)".

2025-03-14 08:54:46 -07:00

partition_spec.py

Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec. We need to do this after sharding-in-types to speed up NamedSharding construction and remove a lot of tech debt and unnecessary complexity.

2025-02-27 08:59:25 -08:00

path.py

Fix detection of epath

2025-03-04 11:44:27 -05:00

pickle_util.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

pjit.py

[explain_cache_miss] Add to explanations the duration of the missed function call

2025-04-14 16:08:24 +03:00

pretty_printer.py

Small cleanup to pretty-printer.

2025-03-24 15:29:29 -04:00

prng.py

remove inaccurate inline comment in PRNGKeyArray constructor

2025-04-15 17:39:40 -07:00

profiler.py

Remove unused return wrapper in annotate_function that creates a self reference cycle loop in python.

2025-04-08 05:25:23 -07:00

public_test_util.py

cleanup now that we depend on ml_dtypes>=0.5

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

random.py

Roll back https://github.com/jax-ml/jax/pull/28022 due to test breakages.

2025-04-15 13:00:04 -07:00

shard_alike.py

Remove internal uses of api_util.shaped_abstractify

2024-12-19 07:06:36 -08:00

sharding_impls.py

Adds a debugging message to assert, otherwise the error is pretty cryptic.

2025-04-14 19:11:02 -07:00

sharding_specs.py

Merge pull request #26518 from superbobry:maint-2

2025-02-13 15:44:19 -08:00

sharding.py

[sharding_in_types] Add vmap + explicit sharding support. The main changes are:

2025-01-29 09:34:27 -08:00

source_info_util.py

[direct-linearize] fix name_stack bugs

2025-03-08 01:51:19 +00:00

sourcemap.py

Run pyupgrade --py310-plus.

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

stages.py

unify stages.Lowering and stages.XlaLowering

2025-04-12 00:31:14 -07:00

test_loader.py

When running test cases concurrently, log the start and end of each test case.

2025-04-10 05:26:02 -07:00

test_util.py

Fix stablehlo version comparison in test utilities.

2025-04-14 13:34:32 -07:00

test_warning_util.py

Switch to a new thread-safe utility for catching warnings.

2025-01-09 11:58:34 -05:00

tpu_custom_call.py

[Pallas][Mosaic][TPU] Add disable_bounds_checks compiler params

2025-04-16 01:01:27 -07:00

traceback_util.py

[cache_misses] Skip tracing-cache-miss explanations for JAX internal functions

2025-04-11 12:53:38 +03:00

tree_util.py

Remove obsolete python key path registry.

2025-04-14 16:33:05 -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

Support __jax_array__ for jnp.shape/jnp.size/jnp.ndim

2025-03-26 15:27:25 -07:00

util.py

Merge tuple_replace and tuple_update in jax._src.util.

2025-04-09 12:50:42 -04:00

xla_bridge.py

Fix stablehlo version comparison in test utilities.

2025-04-14 13:34:32 -07:00

xla_metadata.py

let XLA metadata be unset in nested dynamic scopes

2025-04-01 20:25:19 -07:00