rocm_jax

clusters

Add option to set coordinator lookup timeout for TPU clusters

2024-03-19 20:55:45 -07:00

cudnn

Run pyupgrade --py310-plus.

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

debugger

Merge pull request #20858 from rajasekharporeddy:doc_typos

2024-04-25 06:58:27 -07:00

export

Ported threefry2x32 for GPU to the typed XLA FFI

2024-06-28 06:24:44 -07:00

extend

Add ffi_call function with a similar signature to pure_callback.

2024-07-01 09:40:31 -04:00

image

Run pyupgrade --py310-plus.

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

internal_test_util

Run pyupgrade --py310-plus.

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

interpreters

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

lax

Remove linear parameter from lax.cond_p.

2024-07-01 10:25:42 -04:00

lib

libdevice.10.bc is removed from JAX wheels bundle.

2024-06-27 08:35:59 -07:00

nn

Rename "Example" to "Examples" in docstrings.

2024-06-21 11:43:16 -04:00

numpy

Merge pull request #22167 from rajasekharporeddy:testbranch1

2024-07-01 17:33:37 -07:00

ops

Run pyupgrade --py310-plus.

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

pallas

Removed inspect.Signature hackery from `pl.BlockSpec`

2024-07-02 06:21:59 -07:00

scipy

Fixing the naming conventions in linalg.py

2024-06-27 13:22:39 +05:30

state

Run pyupgrade --py310-plus.

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

third_party

Run pyupgrade --py310-plus.

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

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

abstract_arrays.py

Removed noop # type: ignore comments

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

ad_checkpoint.py

Add a direct HLO lowering of remat_p that doesn't call eval_jaxpr.

2024-06-27 15:15:11 -07:00

ad_util.py

Run pyupgrade --py310-plus.

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

api_util.py

Run pyupgrade --py310-plus.

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

api.py

Run pyupgrade --py310-plus.

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

array.py

Improve error message when a global jax.Array is closed over a jitted function in McJAX.

2024-06-29 14:36:44 -07:00

basearray.py

jnp.take: fix annotation for fill_value

2024-05-25 14:20:55 -07:00

basearray.pyi

Run pyupgrade --py310-plus on .pyi files.

2024-06-26 15:23:57 -04:00

blocked_sampler.py

Run pyupgrade --py310-plus.

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

cache_key.py

Bumped mypy to 1.10.0 and ruff to 0.4.4

2024-05-16 23:16:32 +01:00

callback.py

Add ffi_call function with a similar signature to pure_callback.

2024-07-01 09:40:31 -04:00

checkify.py

Remove linear parameter from lax.cond_p.

2024-07-01 10:25:42 -04:00

cloud_tpu_init.py

Run pyupgrade --py310-plus.

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

compilation_cache_interface.py

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

2024-04-16 20:28:24 -07:00

compilation_cache.py

Implement LRU cache eviction for persistent compilation cache

2024-06-11 21:48:35 +04:00

compiler.py

Run pyupgrade --py310-plus.

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

compute_on.py

Move throwing the error for invalid compute_on into the try-finally context.

2024-05-28 15:20:26 -07:00

config.py

Ported threefry2x32 for GPU to the typed XLA FFI

2024-06-28 06:24:44 -07:00

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

custom_api_util.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

custom_batching.py

Run pyupgrade --py310-plus.

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

custom_derivatives.py

Run pyupgrade --py310-plus.

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

custom_transpose.py

Run pyupgrade --py310-plus.

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

debugging.py

Run pyupgrade --py310-plus.

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

deprecations.py

Register beta args deprecation

2024-06-11 16:19:14 -07:00

dispatch.py

Run single controller device_put via efficient reshard if the device_set of input and the sharding is the same. The transfer_guard in the test fails before this CL.

2024-07-01 13:14:43 -07:00

distributed.py

Rename "Example" to "Examples" in docstrings.

2024-06-21 11:43:16 -04:00

dlpack.py

Removed noop # type: ignore comments

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

dtypes.py

Deprecate passing of arrays in place of dtypes.

2024-06-14 05:40:04 -07:00

earray.py

Make pxla.shard_arg batch calls to xc.copy_array_to_devices_with_sharding

2024-06-13 13:10:10 -07:00

effects.py

Effects is now an immutable set

2023-12-11 08:45:52 -08:00

environment_info.py

jax.print_environment_info: add platform.uname()

2024-02-13 11:23:56 -08:00

errors.py

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07:00

flatten_util.py

Fix duplicate word occurrences

2023-12-19 06:15:30 +01:00

gfile_cache.py

Implement LRU cache eviction for persistent compilation cache

2024-06-11 21:48:35 +04:00

hardware_utils.py

making sure enhanced barrier only turns on when there is a supported TPU available.

2024-01-12 23:47:37 +00: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

Add concrete layout API to JAX. The API takes major_to_minor: tuple[int, ...] and tiling: tuple[tuple[int, ...], ...] as the arguments. Allows users to pass layouts to with_sharding_constraint to constrain the layout + sharding.

2024-06-27 16:47:31 -07:00

lazy_loader.py

Run pyupgrade --py310-plus.

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

linear_util.py

Run pyupgrade --py310-plus.

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

logging_config.py

Add jax_debug_log_modules config option.

2023-07-28 18:11:12 +00:00

lru_cache.py

[LRUCache] Fix CI failure on Windows

2024-06-18 09:46:09 -07:00

maps.py

Add concrete layout API to JAX. The API takes major_to_minor: tuple[int, ...] and tiling: tuple[tuple[int, ...], ...] as the arguments. Allows users to pass layouts to with_sharding_constraint to constrain the layout + sharding.

2024-06-27 16:47:31 -07:00

mesh.py

Rename "Example" to "Examples" in docstrings.

2024-06-21 11:43:16 -04:00

monitoring.py

Upgrade remaining sources to Python 3.9

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

op_shardings.py

Removed noop # type: ignore comments

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

partition_spec.py

Make UNCONSTRAINED's __str__ and __repr__ so it prints nicely in all cases

2024-05-30 17:41:32 -07:00

path.py

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

2024-04-16 20:28:24 -07:00

pickle_util.py

Removed noop # type: ignore comments

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

pjit.py

Add concrete layout API to JAX. The API takes major_to_minor: tuple[int, ...] and tiling: tuple[tuple[int, ...], ...] as the arguments. Allows users to pass layouts to with_sharding_constraint to constrain the layout + sharding.

2024-06-27 16:47:31 -07:00

pretty_printer.py

Changed the naming of internal config APIs

2024-06-18 11:48:57 +01:00

prng.py

Fix error in custom call registration for some FFI functions

2024-06-29 12:18:34 -07:00

profiler.py

Run pyupgrade --py310-plus.

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

public_test_util.py

[Rollback] Bumped the minimum ml_dtypes version to 0.4.0

2024-06-12 14:40:40 -07:00

random.py

Rename "Example" to "Examples" in docstrings.

2024-06-21 11:43:16 -04:00

shard_alike.py

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -07:00

sharding_impls.py

Make device_put work with inputs which are host local and the sharding is global sharding i.e. sharding spanning across multiple hosts.

2024-06-28 09:46:49 -07:00

sharding_specs.py

Remove mesh_shape from ShardingSpec since it's not in use anymore.

2024-04-07 21:13:41 -07:00

sharding.py

Add util.cache to jax.clear_caches and move pjit, sharding, array, etc uses of functools.lru_cache to util.cache so that those caches will be cleared if jax.clear_caches is called.

2024-06-11 12:46:47 -07: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

Run pyupgrade --py310-plus.

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

test_util.py

Run pyupgrade --py310-plus.

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

tpu_custom_call.py

[pallas] Add limited support for shape polymorphism for TPU

2024-07-01 10:18:06 -07:00

traceback_util.py

Run pyupgrade --py310-plus.

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

tree_util.py

Run pyupgrade --py310-plus.

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

tree.py

Run pyupgrade --py310-plus.

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

typing.py

jnp.take: fix annotation for fill_value

2024-05-25 14:20:55 -07:00

util.py

Don't use the trace context in the prune_closed_jaxpr_outputs cache.

2024-07-01 09:41:37 -07:00

xla_bridge.py

Reverts 56526b46aad726ec4632fc3b18c79d26a8f399ef

2024-06-28 12:13:01 -07:00