rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-16 20:06:05 +00:00

History

Peter Hawkins 262a4f482c Deprecate support for custom lowering rules that return tuple-wrapped ir.Values.

https://github.com/google/jax/pull/22211 forbade custom lowering rules from returning singleton tuples of ir.Value, but this appears to break downstream users, notably Transformer Engine. Instead, allow lowering rules to return singleton tuples and unwrap them if needed, but warn if this behavior is seen.

PiperOrigin-RevId: 650345051

2024-07-08 12:54:44 -07:00

clusters

Merge pull request #20174 from coreyjadams:main

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

cudnn

PR #21371 : [NVIDIA] Add new SDPA API to jax.nn

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

debugger

Merge pull request #20858 from rajasekharporeddy:doc_typos

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

export

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

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

Deprecate support for custom lowering rules that return tuple-wrapped ir.Values.

2024-07-08 12:54:44 -07:00

lax

Merge pull request #19614 from cgarciae:batch_map

2024-07-03 15:01:52 -07:00

lib

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

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

fix lint errors at HEAD

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

numpy

Improve documentation of jax.numpy.clip

2024-07-08 06:59:45 -07:00

ops

Run pyupgrade --py310-plus.

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

pallas

Merge pull request #22279 from ayaka14732:ayx/lowering/shift_right_arithmetic

2024-07-08 09:15:27 -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

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

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

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

Added basic support for int2/uint2 dtypes to JAX

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

`debug_info no longer requires non-None func_src_info`

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

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

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

test_util.py

Added some documentation to `jtu.capture_stdout`

2024-07-02 13:46:19 -07:00

tpu_custom_call.py

Introduce a canonicalize pass, rewrite all contractions as matmuls (vector::ContractionOp as tpu::MatMulOp), remove special handling for contraction op in other passes.

2024-07-03 14:49:35 -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

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

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