rocm_jax

clusters

Adds a flag to control proxy env checking.

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

cudnn

Support FP8 for dot_product_attention

2024-12-13 17:56:16 +00:00

debugger

Cleanup: fix unused imports & mark exported names

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

export

[export] Expand exporting to work with AbstractMesh.

2024-12-16 10:30:46 +02:00

extend

Cleanup: remove uses of no-op raise_to_shaped

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

image

Run pyupgrade --py310-plus.

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

internal_test_util

Activate Triangular Solve to XLA's FFI

2024-12-11 02:22:37 -08:00

interpreters

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

lax

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

2024-12-17 10:12:34 -08:00

lib

[Mosaic GPU] Ensure that the dialect module can be loaded successfully.

2024-11-05 00:47:21 -08:00

nn

Add test of relu grad at zero. Update paper links.

2024-12-10 19:39:47 -05:00

numpy

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

2024-12-17 10:12:34 -08:00

ops

Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d

2024-10-04 13:56:27 -07:00

pallas

Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0

2024-12-17 10:12:34 -08:00

scipy

Add Rotation return type hint to Rotation.__mul__()

2024-12-11 12:48:02 +01:00

state

Cleanup: remove uses of no-op raise_to_shaped

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

third_party

feat(lib): add real-valued implementation of jax.scipy.special.fresnel

2024-09-03 09:50:19 +02:00

__init__.py

…

abstract_arrays.py

Cleanup: remove uses of no-op raise_to_shaped

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

ad_checkpoint.py

Fix save_from_both_policies in presence of save_and_offload_only_these_names by comparing the enum

2024-12-16 16:41:13 -08:00

ad_util.py

Cleanup: remove uses of no-op raise_to_shaped

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

api_util.py

Cleanup: remove uses of no-op raise_to_shaped

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

api.py

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

array.py

Don't monkey-patch functions in test_utils to count events for tests.

2024-12-12 09:58:14 -05:00

basearray.py

Make Tracers non-hashable

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

basearray.pyi

Make jnp.subtract a ufunc

2024-10-21 10:11:51 -07:00

blocked_sampler.py

Remove internal KeyArray alias

2024-11-20 10:30:12 -08: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

Fix typos: Change 'arugments' to 'arguments'.

2024-11-13 15:58:45 -05:00

checkify.py

Cleanup: remove uses of no-op raise_to_shaped

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

cloud_tpu_init.py

Fix KeyError recently introduced in cloud_tpu_init.py

2024-11-20 17:46:06 +00:00

compilation_cache_interface.py

Import etils.epath lazily

2024-08-13 14:48:38 -07:00

compilation_cache.py

[AutoPGLE] Explicitly ignore host callback pointers

2024-11-26 04:06:15 -08:00

compiler.py

Remove autotune sharing.

2024-12-13 01:22:27 -08:00

compute_on.py

Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461

2024-11-05 08:32:25 -08:00

config.py

Remove autotune sharing.

2024-12-13 01:22:27 -08:00

core.py

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

custom_api_util.py

…

custom_batching.py

Cleanup: remove uses of no-op raise_to_shaped

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

custom_derivatives.py

Cleanup: remove uses of no-op raise_to_shaped

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

custom_partitioning_sharding_rule.py

[jax:custom_partitioning] Make SdyShardingRule a user facing class.

2024-12-10 13:05:43 -08:00

custom_partitioning.py

Cleanup: remove uses of no-op raise_to_shaped

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

custom_transpose.py

Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.

2024-10-29 11:04:31 -07:00

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

deprecations.py

Log DeprecationWarnings once per method/class

2024-12-14 04:38:53 -08:00

dispatch.py

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

distributed.py

Add an option to deactivate automatic cluster detection in jax.distributed.initialize().

2024-12-02 17:42:12 -08:00

dlpack.py

Deprecated calling `jax.dlpack.from_dlpack` with a DLPack tensor

2024-09-03 15:16:02 -07:00

dtypes.py

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

earray.py

Propagate CopySemantics from python to C++ transfer APIs so that device_put works correctly in presence of copy/donate options that user specified.

2024-11-07 15:51:54 -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

errors.py

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

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

flatten_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

hardware_utils.py

Enable the enhanced TPU launch barrier on all TPU generations.

2024-09-11 15:28:14 -07:00

jaxpr_util.py

Run pyupgrade --py310-plus.

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

lax_reference.py

Implements Ragged Dot API

2024-05-11 06:40:18 -07:00

layout.py

sharding cleanup: use inline checks for unimplemented and auto

2024-10-25 04:22:40 -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

Stop using generators for linear_util transformations.

2024-11-13 13:47:07 -08:00

logging_config.py

Adding JAX_LOGGING_LEVEL configuration option

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

lru_cache.py

Bumped jaxlib version mypy uses on the CI

2024-07-26 11:22:39 +01:00

mesh_utils.py

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

mesh.py

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

monitoring.py

…

op_shardings.py

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

partition_spec.py

[sharding_in_types] Enforce AxisTypes to always exist if set_mesh is used.

2024-12-10 18:03:21 -08:00

path.py

Import etils.epath lazily

2024-08-13 14:48:38 -07:00

pickle_util.py

Removed noop # type: ignore comments

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

pjit.py

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

pretty_printer.py

Bumped mypy to 1.11.0 and jaxlib to 0.4.31 on the CI

2024-08-01 22:30:24 +01:00

prng.py

Lower threefry as an out-of-line MLIR function on TPU.

2024-11-15 08:49:35 -08:00

profiler.py

Small fix for the jax trace dumping path

2024-08-26 10:51:08 -07:00

public_test_util.py

Add float8_e4m3 and float8_e3m4 types support

2024-11-08 18:58:31 +00:00

random.py

fix gamma_p in vmap-based impl rule mode

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

shard_alike.py

[SDY] add JAX lowering to Shardy ShardingGroupOp for shard_alike.

2024-11-08 11:02:50 -08:00

sharding_impls.py

Put abstract_mesh on every eqn so that we can preserve it during eval_jaxpr and check_jaxpr roundtrip.

2024-12-17 09:17:21 -08:00

sharding_specs.py

Fix type annotations for NumPy 2.2

2024-12-11 14:24:58 -08:00

sharding.py

[jax] Typing on common_devices_indices_map

2024-12-02 12:40:36 -08:00

source_info_util.py

Use a class with __slots__ instead of a NamedTuple in JaxprEqn and SourceInfo, which are two tuples we build frequently.

2024-07-01 19:18:58 -07:00

sourcemap.py

Run pyupgrade --py310-plus.

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

stages.py

Don't monkey-patch functions in test_utils to count events for tests.

2024-12-12 09:58:14 -05:00

test_util.py

Merge pull request #25426 from hawkinsp:tls

2024-12-12 08:15:34 -08:00

tpu_custom_call.py

jax.core: deprecate a number of APIs

2024-12-10 11:11:32 -08:00

traceback_util.py

Run pyupgrade --py310-plus.

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

tree_util.py

Implement flatten one level with keys in C++ and use it for the prefix/equality error printing.

2024-12-12 16:37:09 -08: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

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

util.py

Don't monkey-patch functions in test_utils to count events for tests.

2024-12-12 09:58:14 -05:00

xla_bridge.py

Add an option to specify mock GPU topology

2024-11-12 08:36:27 -08:00

xla_metadata.py

Reverts a913fbf2fddc5b8c1b6c85b159d0eeb1bf65d461

2024-11-05 08:32:25 -08:00