rocm_jax

clusters

Adds a flag to control proxy env checking.

2024-11-16 00:00:59 -08:00

cudnn

Support nvfp4

2025-03-11 19:33:25 +00:00

debugger

Cleanup: fix unused imports & mark exported names

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

export

Make the signature of AbstractMesh to be AbstractMesh(axis_size: tuple[int, ...], axis_name: tuple[str, ...], *, axis_types) instead of AbstractMesh(shape_tuple: tuple[tuple[str, int], ...], *, axis_types) so that we are consistent across all Mesh APIs: Mesh, AbstractMesh and make_mesh

2025-03-12 21:32:31 -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

#sdy Add JAX backwards compatibility test.

2025-03-03 06:01:34 -08:00

interpreters

Make sure to DCE read effects

2025-03-18 18:42:14 -07:00

lax

fix to ragged_all_to_all transpose

2025-03-18 12:51:21 -07:00

lib

[Mosaic GPU] Remove old jaxlib version guards.

2025-02-12 08:49:40 -08:00

nn

Merge pull request #26765 from Qwlouse:patch-1

2025-03-04 08:30:45 -08:00

numpy

Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/

2025-03-18 16:29:37 -07:00

ops

doc: fix description of logsumexp axis

2025-02-28 12:53:33 -08:00

pallas

Merge pull request #27227 from jburnim:jburnim_pallas_interpret_mode4

2025-03-18 20:22:27 -07:00

scipy

Fix scipy.special.gammainc/gammaincc evaluation at boundary points

2025-03-11 21:18:47 +02:00

state

Add a variant of safe_map() that has no return value, named foreach().

2025-03-14 07:42:48 -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

Rename dot_with_no_batch_dims_saveable to dots_with_no_batch_dims_saveable for internal consistency

2025-03-10 13:04:49 -07:00

ad_util.py

Cleanup: remove uses of no-op raise_to_shaped

2024-12-12 09:49:06 -08:00

api_util.py

Internal: make it easier to detect the vmap sentinel

2025-03-10 11:37:50 -07:00

api.py

Allow pspec to be passed to device_put if there is a mesh in the surrounding context

2025-03-17 17:47:56 -07:00

array.py

Check the type of mesh in use_abstract_mesh and use_concrete_mesh

2025-03-18 16:57:40 -07:00

basearray.py

Make Tracers non-hashable

2024-11-05 09:08:33 -08:00

basearray.pyi

[sharding_in_types] Rename .at[...].get(out_spec) to .at[...].get(out_sharding).

2025-01-16 18:56:52 -08: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

Remove dead code after minimum jaxlib version bump to v0.4.36.

2024-12-09 07:35:05 -08:00

callback.py

Reverts 6e83de5909b5dad6827c1682726c840c9688b32d

2025-02-21 18:34:16 -08:00

checkify.py

Add a variant of safe_map() that has no return value, named foreach().

2025-03-14 07:42:48 -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

Add JAX_COMPILATION_CACHE_EXPECT_PGLE option

2025-02-06 08:19:45 +00:00

compiler.py

Remove code present to support jaxlib < 0.5.1.

2025-02-26 07:40:40 -05:00

compute_on.py

Minor format cleanup

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

config.py

[XLA:GPU] Add missing BF16_BF16_F32_X9 matmul option in config.py

2025-03-14 08:58:59 -07:00

core.py

Rename AxisTypes enum to AxisType

2025-03-14 11:48:21 -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

Reverts 342cb7b99a09180472823a33c7cdad8a8db77875

2025-03-05 10:22:40 -08: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

Support tuples in custom_partitioning.

2025-03-18 14:57:08 -07:00

custom_transpose.py

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

2025-02-08 09:13:55 +02:00

debugging.py

Replace cached function get_replicated_hlo_sharding() with a constant.

2025-03-17 13:17:33 -07:00

deprecations.py

Add sph_harm_y to jax.scipy.special and deprecate sph_harm.

2025-01-09 12:53:00 -05:00

dispatch.py

Enable jax.device_put to a sharding with no local devices.

2025-03-17 16:49:46 -07:00

distributed.py

add distributed.is_initialized

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

dlpack.py

Pulled DLDeviceType to XLA backend mapping into a global

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

dtypes.py

Fix lint

2025-03-07 04:52:58 +00: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

…

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

Support error checking in explicit mode

2025-03-14 18:58:26 -07:00

errors.py

Consolidate material on PRNGs and add a short summary to Key Concepts.

2024-11-15 14:44:57 -08:00

ffi.py

Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.

2025-02-21 09:45:59 -08:00

flatten_util.py

Avoid calls to warnings.catch_warnings in JAX core code.

2024-12-20 15:43:03 -05: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

sharding cleanup: use inline checks for unimplemented and auto

2024-10-25 04:22:40 -07:00

lazy_loader.py

…

linear_util.py

[better_errors] Port the Pallas debug info mechanisms to the new JAX DebugInfo.

2025-02-25 14:43:17 +01: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

Reverts 0f0d5e90ef1c3d60f35020141710ea350d17816b

2025-02-27 13:27:32 -08:00

mesh.py

Check the type of mesh in use_abstract_mesh and use_concrete_mesh

2025-03-18 16:57:40 -07:00

monitoring.py

Add JAX events that have time spans, not only durations.

2025-01-07 23:08:30 -08:00

named_sharding.py

Rename AxisTypes enum to AxisType

2025-03-14 11:48:21 -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

…

pjit.py

Rename AxisTypes enum to AxisType

2025-03-14 11:48:21 -07:00

pretty_printer.py

Make api_test.py work when test cases are run using multiple threads.

2025-01-08 14:09:07 -08:00

prng.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

profiler.py

Add JAX_COMPILATION_CACHE_EXPECT_PGLE option

2025-02-06 08:19:45 +00:00

public_test_util.py

support e2m1fn

2025-03-05 17:44:34 +00:00

random.py

Add replace option to random.categorical to enable sampling without replacement.

2025-03-17 13:41:46 -04:00

shard_alike.py

Remove internal uses of api_util.shaped_abstractify

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

sharding_impls.py

Check the type of mesh in use_abstract_mesh and use_concrete_mesh

2025-03-18 16:57:40 -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

…

stages.py

Remove obsolete fallback for cost analysis.

2025-03-13 14:49:40 -07:00

test_util.py

Rename AxisTypes enum to AxisType

2025-03-14 11:48:21 -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

[JAX/Pallas] Add has_side_effect parameter to CompilerParams to stop CSE of operations.

2025-02-11 13:33:01 -08:00

traceback_util.py

…

tree_util.py

Removed unused `# type: ignore` comments

2025-02-13 21:12:27 +00: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

…

util.py

Remove code that preserved _original_py_fns on C++ classes.

2025-03-17 12:43:04 -07:00

xla_bridge.py

Raise an error if jax.config.update('jax_num_cpu_devices', val) is called after backend is initialized

2025-03-13 14:45:53 -07:00

xla_metadata.py

Fix CI failures in xla_metadata_test.

2025-01-30 17:28:54 -08:00