rocm_jax

clusters

Upgrade remaining sources to Python 3.9

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

cudnn

typing: fix incorrect tuple annotations

2024-02-26 10:53:19 -08:00

debugger

jax.debug.breakpoint now supports being DCE'd. Drive-by: fix #16186

2023-07-31 15:36:43 -07:00

extend

Upgrade most .py sources to 3.9

2023-12-08 12:23:15 +00:00

image

Upgrade remaining sources to Python 3.9

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

internal_test_util

[JAX] Add an option subset_by_index that allows computing a contiguous subset of singular components from svd.

2024-02-15 16:33:09 -08:00

interpreters

[XLA:MHLO->HLO] Allow partially-set parameter tuple sharding to exist by filling in the missing sharding elements with replicated sharding. (This is what is done for the missing shardings in the result tuple.)

2024-03-14 15:09:56 -07:00

lax

Reverts 8e2a8b7b95e838947dcf581d146909d5c4128742

2024-03-13 07:01:49 -07:00

lib

Handle TypeError due to | in type annotations in Triton MLIR bindings

2024-02-22 02:59:46 -08:00

nn

Merge pull request #17741 from froystig:new-style-key-docs

2024-03-08 16:41:22 -08:00

numpy

Repair various type errors

2024-03-13 15:13:56 -04:00

ops

jax.numpy: implement scalar boolean indexing

2024-02-09 11:00:00 -08:00

pallas

[Pallas] Raise NotImplementedError for strided load/store in interpret mode.

2024-03-14 19:44:39 -07:00

scipy

DOC: fix two minor doc issues

2024-03-05 15:22:42 -08:00

state

[Pallas] Raise NotImplementedError for strided load/store in interpret mode.

2024-03-14 19:44:39 -07:00

third_party

Rename _wraps to implements

2024-01-24 14:14:19 -08: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

Remove code that guards against int4 not being available.

2024-02-15 08:41:17 -08:00

ad_checkpoint.py

ignore NamedAxisEffect for remat and dce purposes

2024-03-05 11:04:23 -08: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

typing: fix incorrect tuple annotations

2024-02-26 10:53:19 -08:00

api.py

[JAX] Convert stablehlo to MLIR bytecode, not an MLIR string.

2024-03-13 06:02:18 -07:00

array.py

Update

2024-03-12 12:56:22 +00:00

basearray.py

Upgrade remaining sources to Python 3.9

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

basearray.pyi

Add copy argument to Array.__array__

2024-03-05 09:31:16 -08:00

cache_key.py

Enhance compilation cache key generation with a custom hook.

2024-02-26 18:18:35 -08:00

callback.py

Log the exception if the callback passed to jax.*_callback raises

2024-03-13 07:23:29 -07:00

checkify.py

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00

cloud_tpu_init.py

Increase minimum jaxlib version to 0.4.20.

2024-02-21 12:58:57 -08:00

compilation_cache_interface.py

Replaced most usages of abc.ABC with util.StrictABC

2024-01-29 12:40:43 +00:00

compilation_cache.py

Increase some log levels to info in public methods.

2024-02-23 08:38:33 -08:00

compiler.py

Add debug logging for autotune profile sharing.

2024-03-11 22:46:56 -07:00

config.py

Reverts 8e2a8b7b95e838947dcf581d146909d5c4128742

2024-03-13 07:01:49 -07:00

core.py

Merge pull request #20218 from mattjj:mutable-arrays-closure

2024-03-13 10:23:23 -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

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00

custom_derivatives.py

temporarily relax the cotangent dtype check introduced in #19009

2024-03-14 13:22:42 -07:00

custom_transpose.py

Upgrade remaining sources to Python 3.9

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

debugging.py

Cache ClosedJaxpr creation to minimize cache_misses. ClosedJaxpr should always be created under a cache.

2023-12-21 22:15:52 -08:00

deprecations.py

Add (private) mechanism for registering and accelerating deprecations

2024-02-07 12:28:03 -08:00

dispatch.py

Allow sharding propagation to input for prng keys whose sharding is not specified.

2024-02-28 15:22:16 -08:00

distributed.py

Upgrade remaining sources to Python 3.9

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

dlpack.py

Fix unnecessary memory copies between GPU and CPU when jax2tf.call_tf() is used.

2024-02-27 10:35:50 -08:00

dtypes.py

typing: fix incorrect tuple annotations

2024-02-26 10:53:19 -08: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

Fix doc typos

2024-03-13 14:48:31 +05:30

flatten_util.py

Fix duplicate word occurrences

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

gfile_cache.py

Fix FileExistsError for tests/gfile_cache_test.py::FileSystemCacheTest::test_threads

2023-06-14 08:07:22 +08: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

Upgrade remaining sources to Python 3.9

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

lax_reference.py

test: fix dtype in test of lax.logistic

2024-01-16 11:59:30 -08:00

layout.py

Make the SpecifiedLayout class opaque.

2023-11-18 15:17:16 -08:00

lazy_loader.py

Apply pyupgrade --py39-plus.

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

linear_util.py

Avoid a couple more uses of operator.attrgetter(), since it's unnecessarily slow compared to writing a function directly.

2024-02-23 09:20:56 -08:00

logging_config.py

Add jax_debug_log_modules config option.

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

maps.py

Add a deprecation warning to the xmap function and the corresponding tutorial

2024-03-11 21:31:32 +00:00

mesh.py

Add __str__ to Mesh so that in jaxprs the mesh doesn't print all the device ids.

2024-01-18 11:23:25 -08:00

monitoring.py

Upgrade remaining sources to Python 3.9

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

op_shardings.py

Allow sharding propagation to input for prng keys whose sharding is not specified.

2024-02-28 15:22:16 -08:00

partition_spec.py

Improve the sharding documentation.

2023-08-03 11:12:02 -04:00

path.py

Make the etils[epath] dependency of JAX optional.

2022-10-12 21:16:48 +00:00

pickle_util.py

Use cloudpickle for serializing host callbacks

2023-06-09 13:34:43 -07:00

pjit.py

[mutable-arrays] support closed-over mutable arrays in jit

2024-03-13 09:59:03 -07:00

pretty_printer.py

Replaced most usages of abc.ABC with util.StrictABC

2024-01-29 12:40:43 +00:00

prng.py

[key reuse] add jax.random.clone

2024-03-08 09:06:00 -08:00

profiler.py

Upgrade remaining sources to Python 3.9

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

public_test_util.py

Support float8 in reduce_min & reduce_max

2023-12-18 13:37:45 -08:00

random.py

Merge pull request #17741 from froystig:new-style-key-docs

2024-03-08 16:41:22 -08:00

shard_alike.py

Guard shard_alike usage on xla_extension_version

2024-01-19 13:02:29 -08:00

sharding_impls.py

[XLA:Python] Port sharding and device lists to nanobind.

2024-03-08 07:54:55 -08:00

sharding_specs.py

Add a new experimental option jax_pmap_no_rank_reduction.

2024-01-19 03:53:37 -08:00

sharding.py

Remove indices and devices from shard_arg_handlers and shard_args.

2024-01-05 14:17:14 -08:00

source_info_util.py

Speed up user filename identification.

2024-02-13 11:04:47 -08:00

stages.py

[JAX] Convert stablehlo to MLIR bytecode, not an MLIR string.

2024-03-13 06:02:18 -07:00

test_util.py

Merge pull request #20128 from shuhand0:dev/shuhan/ci2

2024-03-12 17:45:19 -07:00

tpu_custom_call.py

A fusion flag for each operand is set to false by default. A custom call writer is expected to turn them on if he expects those fusions to be profitable. The operand may not fuse despite the flag being turned to true because of other constrains such as estimated memory required after fusion.

2024-03-11 13:21:06 -07:00

traceback_util.py

Upgrade remaining sources to Python 3.9

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

tree_util.py

tree_transpose: optionally infer inner_treedef

2024-02-15 12:01:21 -08:00

tree.py

Add jax.tree module with aliases of jax.tree_util

2024-02-12 13:07:59 -08:00

typing.py

[shape_poly] Simplify the API for processing polymorphic_shape specifications

2023-11-28 12:45:59 +02:00

util.py

Repair various type errors

2024-03-13 15:13:56 -04:00

xla_bridge.py

Update

2024-03-03 19:57:26 +00:00