rocm_jax

clusters

Merge pull request #20174 from coreyjadams:main

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

cudnn

Fixed rules where `sliding_window_length` was not forwarded

2024-09-03 06:24:01 -07:00

debugger

Merge pull request #20858 from rajasekharporeddy:doc_typos

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

export

Fixed some type errors under pyright

2024-09-05 09:56:38 +01:00

extend

Add support for the DeviceLocalLayout API when lowering FFI calls.

2024-09-05 14:30:06 -04:00

image

Run pyupgrade --py310-plus.

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

internal_test_util

Activate Householder Product to XLA's FFI

2024-09-02 06:19:01 -07:00

interpreters

Fixed some type errors under pyright

2024-09-05 09:56:38 +01:00

lax

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

lib

Partially rolling forward #22998

2024-09-02 04:44:47 -07:00

nn

Merge pull request #22516 from kaixih:support_variable_seqlen

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

numpy

Refactor & document cumulative reductions

2024-09-04 12:20:24 -07:00

ops

Reverts 6fc57c0eb6f06b2da20c94f5f127fe4a551bda09

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

pallas

[Mosaic GPU] Remove the unnecessary scratch space operand

2024-09-05 04:57:52 -07:00

scipy

Fixed some type errors under pyright

2024-09-05 09:56:38 +01:00

state

Support strided load / store in interpret mode

2024-08-28 05:10:54 -07:00

third_party

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

2024-09-03 09:50:19 +02: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 various ununsed functions

2024-08-01 11:18:19 +01:00

ad_checkpoint.py

Merge pull request #22263 from hawkinsp:tuples

2024-07-17 09:56:18 -07:00

ad_util.py

Run pyupgrade --py310-plus.

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

api_util.py

Consolidate handling of input argument resolution in custom_* APIs.

2024-08-13 00:30:23 -07:00

api.py

Reverts 82c9da020a78997862a8f7ccd494bed363f7ed01

2024-08-29 09:43:19 -07:00

array.py

Initial integration of sharding in types in JAX. Currently we just support nary ops in forward only sharding propagation. Currently this functionality is experimental and hidden behind jax_sharding_in_types config flag.

2024-08-29 10:50:04 -07:00

basearray.py

Document copy_to_host_async method of jax.Array

2024-07-26 08:21:01 -07:00

basearray.pyi

Fix pytype errors and args for jax.Array methods

2024-09-03 17:06:45 +00: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

Set the value of config.jax_platforms directly instead of setting the env variable JAX_PLATFORMS. Setting JAX_PLATFORMS doesn’t do anything because config.jax_platforms is already initialized at that point.

2024-07-22 15:09:15 -07:00

compilation_cache_interface.py

Import etils.epath lazily

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

compilation_cache.py

Merge pull request #22556 from cool-RR:log-cache-key

2024-07-26 05:32:11 -07:00

compiler.py

[Pallas][XLA:Mosaic] Add python stack traces to Mosaic errors that occur in Pallas.

2024-08-12 14:42:48 -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

Add use_shardy_partitioner to thread local jit state

2024-09-02 08:47:06 -07:00

core.py

Remove dtype and weak_type from __slots__ of ShapedArray since it comes from UnShapedArray

2024-08-30 14:37:29 -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

Consolidate handling of input argument resolution in custom_* APIs.

2024-08-13 00:30:23 -07:00

custom_derivatives.py

Document methods of custom_jvp/custom_vjp

2024-08-14 15:37:20 -07:00

custom_partitioning.py

Make custom partitioning work without a mesh context manager. If the arguments have NamedSharding on them, then inside partition function, we should get NamedSharding without the existence of the mesh context manager

2024-08-12 10:40:31 -07:00

custom_transpose.py

Run pyupgrade --py310-plus.

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

debugging.py

Added a trivial discharge rule for debug_callback_p

2024-07-29 22:26:01 +01:00

deprecations.py

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

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

dispatch.py

If input layouts are specified via in_shardings to jit and the array that the jitted function is called with is uncommitted, reshard the input array to the layout specified by the user.

2024-08-19 15:10:32 -07:00

distributed.py

Register shutdown code at import to hopefully get registered before any other atexit callbacks.

2024-08-12 11:29:08 -07:00

dlpack.py

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

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

dtypes.py

Fix a deprecation warning for NumPy array conversion.

2024-09-04 12:41:28 -07:00

earray.py

If input layouts are specified via in_shardings to jit and the array that the jitted function is called with is uncommitted, reshard the input array to the layout specified by the user.

2024-08-19 15:10:32 -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

Add TPU v6e to the list of known TPUs.

2024-09-03 14:20:35 -04: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

If input layouts are specified via in_shardings to jit and the array that the jitted function is called with is uncommitted, reshard the input array to the layout specified by the user.

2024-08-19 15:10:32 -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

Bumped mypy to 1.11.0 and jaxlib to 0.4.31 on the CI

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

logging_config.py

Add jax_debug_log_modules config option.

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

lru_cache.py

Bumped jaxlib version mypy uses on the CI

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

mesh_utils.py

Create jax.make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str], devices: Sequence[jax.Device] | None = None) API to make it easier to create a mesh and reduce a ton of boilerplate.

2024-09-03 14:32:03 -07:00

mesh.py

Fix empty mesh size and abstract_mesh

2024-08-20 10:00:37 -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

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

Reverts 82c9da020a78997862a8f7ccd494bed363f7ed01

2024-08-29 09:43:19 -07: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

If input layouts are specified via in_shardings to jit and the array that the jitted function is called with is uncommitted, reshard the input array to the layout specified by the user.

2024-08-19 15:10:32 -07:00

profiler.py

Small fix for the jax trace dumping path

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

public_test_util.py

Added basic support for int2/uint2 dtypes to JAX

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

random.py

Merge pull request #23257 from carlosgmartin:random-orthogonal-citation

2024-09-04 17:24:51 -07:00

shard_alike.py

Bump minimum jaxlib version to 0.4.27

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

sharding_impls.py

Split physical axes by default when device kind is TPU v5 lite to allow for mesh shapes (2, 2) when there are 8 v5e devices on a 4x2 topology.

2024-09-04 11:49:17 -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 num_devices to Sharding interface so that it works with NamedSharding containing AbstractMesh too.

2024-08-14 09:03:17 -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

Make input_layouts and output_layouts properties of Compiled. This is equivalent to input_shardings and output_shardings which are also properties on Compiled.

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

test_util.py

Rename jtu.create_global_mesh to jtu.create_mesh and use jax.make_mesh inside jtu.create_mesh to get maximum test coverage of the new API.

2024-09-03 16:23:07 -07:00

tpu_custom_call.py

Allow optionally coloring output arguments from tpu_custom_call

2024-09-04 16:33:25 -07:00

traceback_util.py

Run pyupgrade --py310-plus.

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

tree_util.py

`jax.tree_util.register_dataclass now validates data_fields and meta_fields`

2024-08-30 02:01:50 -07:00

tree.py

Fix annotation for jax.tree.transpose

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

typing.py

[jax] Make the SupportsDType protocol runtime checkable.

2024-07-16 10:57:15 -07:00

util.py

Collapsed a few unnecessary `if TYPE_CHECKING` blocks

2024-08-12 13:08:55 +01:00

xla_bridge.py

Add a jax.process_indices function.

2024-08-12 10:30:41 -07:00