rocm_jax

clusters

Merge pull request #20174 from coreyjadams:main

2024-07-08 12:19:18 -07:00

cudnn

Raise NotImplementedError instead of assert for unsupported Q dtype in fused attention.

2024-07-09 07:37:13 -07:00

debugger

Merge pull request #20858 from rajasekharporeddy:doc_typos

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

export

[export] Fix poly shape check for vjp function with integer valued, polymorphic output.

2024-07-10 06:12:19 -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

Add sharding to convert_element_type_p primitive.

2024-07-09 07:33:29 -07:00

interpreters

Add sharding to convert_element_type_p primitive.

2024-07-09 07:33:29 -07:00

lax

Deduped three identical implementations of `hoist_consts_to_refs`

2024-07-12 12:48:28 +01:00

lib

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

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

nn

Fix documentation for dot_product_attention.

2024-07-09 08:09:15 -07:00

numpy

Merge pull request #22347 from rajasekharporeddy:testbranch3

2024-07-10 10:18:05 -07:00

ops

Run pyupgrade --py310-plus.

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

pallas

Deduped three identical implementations of `hoist_consts_to_refs`

2024-07-12 12:48:28 +01:00

scipy

jnp.vectorize: respect numpy_rank_promotion config

2024-07-08 09:03:03 -07:00

state

Deduped three identical implementations of `hoist_consts_to_refs`

2024-07-12 12:48:28 +01: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

Added basic support for int2/uint2 dtypes to JAX

2024-07-04 04:13:24 -07:00

ad_checkpoint.py

Fix lint error at head.

2024-07-03 15:06:23 -04:00

ad_util.py

Run pyupgrade --py310-plus.

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

api_util.py

`debug_info no longer requires non-None func_src_info`

2024-07-05 20:08:53 +01:00

api.py

Update jax/_src/api.py

2024-07-02 16:10:36 +01:00

array.py

Don't wrap singleton ir.Values with tuples during HLO lowering.

2024-07-01 16:11:00 -04: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

Merge pull request #22211 from hawkinsp:singletons

2024-07-03 11:07:00 -07:00

checkify.py

Remove linear parameter from lax.cond_p.

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

cloud_tpu_init.py

Enable runtime uptime telemetry for JAX on Cloud TPU, only if the flag is not set by the user explicitly. Otherwise, prefer the user preference.

2024-07-02 12:48:48 -07:00

compilation_cache_interface.py

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

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

compilation_cache.py

Clean up

2024-07-04 17:16:32 +04:00

compiler.py

Clean up

2024-07-04 17:16:32 +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

`debug_info no longer requires non-None func_src_info`

2024-07-05 20:08:53 +01: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

Don't wrap singleton ir.Values with tuples during HLO lowering.

2024-07-01 16:11:00 -04:00

custom_partitioning.py

Moved the implementation of `custom_partitioning` into jax/_src

2024-07-08 04:31:44 -07:00

custom_transpose.py

Run pyupgrade --py310-plus.

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

debugging.py

Don't wrap singleton ir.Values with tuples during HLO lowering.

2024-07-01 16:11:00 -04:00

deprecations.py

Register beta args deprecation

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

dispatch.py

[JAX] Check if an array is deleted when resharding it with a different device order

2024-07-10 12:57:06 -07:00

distributed.py

Squashed commit of the following:

2024-07-02 13:18:05 -05:00

dlpack.py

Removed noop # type: ignore comments

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

dtypes.py

extend type annotation for lax.convert_element_type

2024-07-11 05:27:11 -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

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

`debug_info no longer requires non-None func_src_info`

2024-07-05 20:08:53 +01:00

logging_config.py

Add jax_debug_log_modules config option.

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

lru_cache.py

Clean up

2024-07-04 17:16:32 +04:00

maps.py

Don't wrap singleton ir.Values with tuples during HLO lowering.

2024-07-01 16:11:00 -04: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

Clean up

2024-07-04 17:16:32 +04:00

pickle_util.py

Removed noop # type: ignore comments

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

pjit.py

Add sharding to convert_element_type_p primitive.

2024-07-09 07:33:29 -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

Added basic support for int2/uint2 dtypes to JAX

2024-07-04 04:13:24 -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

require keyword arguments to Traced.lower

2024-07-10 17:59:26 -07:00

test_util.py

Added test assertions for `pl.debug_print` on TPU

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

tpu_custom_call.py

Use default tiling in scratch buffers if XLA enables it

2024-07-08 22:49:10 -07:00

traceback_util.py

Run pyupgrade --py310-plus.

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

tree_util.py

[pallas] Improve some error messages and add API tests.

2024-07-04 09:02:16 +02:00

tree.py

Fix annotation for jax.tree.transpose

2024-07-09 08:06:16 -07:00

typing.py

jnp.take: fix annotation for fill_value

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

util.py

Prefer __qualname__ as a pjit_p name.

2024-07-02 14:53:50 -04:00

xla_bridge.py

Reverts 56526b46aad726ec4632fc3b18c79d26a8f399ef

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