Yash Katariya b1f7627c71 [Rollback] Bumped the minimum ml_dtypes version to 0.4.0 ...

Reverts e86c436e7f8e4e0546eff8bc2d3756a7c49dc83b

PiperOrigin-RevId: 642741832

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

..

clusters

Add option to set coordinator lookup timeout for TPU clusters

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

cudnn

init

2024-05-24 22:26:20 +00:00

debugger

Merge pull request #20858 from rajasekharporeddy:doc_typos

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

export

Document jax.export serialization version numbers

2024-06-12 12:44:42 -07:00

extend

Edit pycapsule docstring to provide a little bit more context

2024-06-07 13:07:03 -04:00

image

Upgrade remaining sources to Python 3.9

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

internal_test_util

[export] Create the jax.export module APIs.

2024-06-10 19:31:51 +02:00

interpreters

[pallas] Add support for cross-platform lowering

2024-06-12 08:48:58 +02:00

lax

DOC: jax.lax.top_k: fix docstring rendering & add example

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

lib

Make CuDNN SDPA API work with JAX with a CUDA plugin configuration.

2024-06-06 12:09:19 -04:00

nn

softmax: document NaN outputs for infinite inputs

2024-05-29 15:00:20 -07:00

numpy

Remove unused variable

2024-06-10 16:30:44 -07:00

ops

Support scalar boolean indices in arr.at[idx].set(vals)

2024-05-20 05:33:36 -07:00

pallas

[Pallas] Allow keys as input to Pallas kernels.

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

scipy

Register beta args deprecation

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

state

Slightly rearranged NDIndexer.from_indices_shape and added missing tests

2024-06-03 12:33:14 +01:00

third_party

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01: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

Initialize JaxprEqnContext only in new_jaxpr_eqn and new_eqn_recipe with the current active compute type if no ctx is specified.

2024-05-22 15:16:58 -07:00

ad_util.py

partial rollback of #19096 due to internal breakage (relying on jax internals)

2023-12-22 15:54:26 -08:00

api_util.py

Removed noop # type: ignore comments

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

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

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

basearray.py

jnp.take: fix annotation for fill_value

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

basearray.pyi

Fix type annotation for index-update

2024-05-28 07:15:15 -07: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

Clarified the type of the inputs to callback APIs

2024-05-16 11:29:09 -07:00

checkify.py

Removed noop # type: ignore comments

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

cloud_tpu_init.py

Enable runtime uptime telemetry for JAX on Cloud TPU.

2024-06-12 13:36:00 -07: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

Allow disabling compilation cache for particular runtime_types.

2024-06-04 13:30:42 -07: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

Implement LRU cache eviction for persistent compilation cache

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

core.py

fix error message for vjp arguments

2024-05-30 21:22:35 +00: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

Removed noop # type: ignore comments

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

custom_derivatives.py

[custom_vjp] allow bwd rule to produce top-level list (not just tuple)

2024-05-31 21:49:06 +00:00

custom_transpose.py

Removed noop # type: ignore comments

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

debugging.py

Merge pull request #21261 from superbobry:mypy-ruff

2024-05-17 00:10:27 -07:00

deprecations.py

Register beta args deprecation

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

dispatch.py

Deprecate XLACompatibleSharding in favor of jax.sharding.Sharding.

2024-06-05 09:07:27 -07:00

distributed.py

Fixed the logging due to a recent change.

2024-06-01 14:28:55 -07:00

dlpack.py

Removed noop # type: ignore comments

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

dtypes.py

[Rollback] Bumped the minimum ml_dtypes version to 0.4.0

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

earray.py

Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.

2024-06-03 14:52:50 -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

Removed noop # type: ignore comments

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

lax_reference.py

Implements Ragged Dot API

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

layout.py

Prefer raising of TypeError for invalid types instead of ValueError.

2024-04-08 13:08:24 +05:30

lazy_loader.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04:00

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

logging_config.py

Add jax_debug_log_modules config option.

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

lru_cache.py

Implement LRU cache eviction for persistent compilation cache

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

maps.py

Remove jaxpr and name from Lowered because specialize already has those. This keeps the abstraction boundary clear. Adapt export to use specialize.

2024-06-06 11:38:56 -07:00

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

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

pretty_printer.py

Compute source maps when pretty-printing jaxprs.

2024-05-06 15:45:25 -04:00

prng.py

Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.

2024-06-03 14:52:50 -07:00

profiler.py

[JAX] Automatically share PGO data for GPU latency-hiding scheduler.

2024-05-29 01:50:03 -07: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

fp8 matmul in pallas

2024-06-07 08:17:06 -07:00

shard_alike.py

Bump minimum jaxlib version to 0.4.27

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

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

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

Removed noop # type: ignore comments

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

sourcemap.py

Merge pull request #21096 from gspschmid:gschmid/sourcemaps

2024-05-08 05:44:08 -07:00

stages.py

[export] Create the jax.export module APIs.

2024-06-10 19:31:51 +02:00

test_util.py

test cleanup: use ExitStack to reduce test boilerplate

2024-06-06 14:18:27 -07:00

tpu_custom_call.py

Removed kernel_regeneration_util from Mosaic

2024-06-11 02:36:41 -07:00

traceback_util.py

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

tree_util.py

tree_all: add support for is_leaf

2024-06-10 09:46:15 -07:00

tree.py

tree_all: add support for is_leaf

2024-06-10 09:46:15 -07:00

typing.py

jnp.take: fix annotation for fill_value

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

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

xla_bridge.py

Enable jax's cloud-tpu configs when libtpu is present via through "pip install" or set by custom through the $TPU_LIBRARY_PATH env var

2024-06-12 11:55:43 -07:00