Yash Katariya 84156f359f Add identity jit tests to go from pinned_host -> device and vice versa ...

PiperOrigin-RevId: 620114420

2024-03-28 18:20:32 -07:00

..

clusters

Add option to set coordinator lookup timeout for TPU clusters

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

cudnn

add is_training && fix seqlen/head_dim checks

2024-03-14 14:34:40 -07: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

Add identity jit tests to go from pinned_host -> device and vice versa

2024-03-28 18:20:32 -07:00

lax

Allows for splitting the transpose of a scan into a scan and a map.

2024-03-28 10:54:50 -07:00

lib

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

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

nn

Require array-like inputs to sparse_plus

2024-03-19 09:06:18 -07:00

numpy

Remove _maybe_device_put because jax.device_put accepts None on the device parameter

2024-03-22 10:39:57 -07:00

ops

jax.numpy: implement scalar boolean indexing

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

pallas

Allows for splitting the transpose of a scan into a scan and a map.

2024-03-28 10:54:50 -07:00

scipy

Add Hilbert matrix to jax.scipy.linalg

2024-03-20 22:55:03 +05:30

state

[Pallas TPU] Add support for dynamic sized (tile aligned) DMAs

2024-03-22 16:59:32 -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

make print_saved_residuals show more checkpoint_name names

2024-03-17 17:55:29 -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

Add a fast path for Python scalars to shaped_abstractify.

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

api.py

Refactorings to the jit implementation.

2024-03-21 05:37:32 -07:00

array.py

[Take 2] Expose .layout on jax.Array. Also add checks in the AOT path to make sure that the input Array's layout matches the layout given to jax.jit.

2024-03-25 10:08:43 -07: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

Strip device_assignment on GPU platform.

2024-03-19 01:50:39 -07:00

callback.py

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

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

checkify.py

Allows for splitting the transpose of a scan into a scan and a map.

2024-03-28 10:54:50 -07: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

[key reuse] rename flag to jax_debug_key_reuse

2024-03-22 05:37:30 -07:00

core.py

Reverts 0dde8f7f9607d09841ece7125dfc0773c3613fab

2024-03-26 22:26:41 -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

Enable remat lowering and tests on GPU

2024-03-15 10:19:53 -07:00

distributed.py

Add option to set coordinator lookup timeout for TPU clusters

2024-03-19 20:55:45 -07: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

Add a fast path for Python scalars to shaped_abstractify.

2024-03-21 18:05: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

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

[Take 2] Expose .layout on jax.Array. Also add checks in the AOT path to make sure that the input Array's layout matches the layout given to jax.jit.

2024-03-25 10:08:43 -07: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

Remove the canonicalization to GSPMDSharding internally in jit. This is not required anymore since the caches are split into tracing, lowering and compilation.

2024-03-26 13:28:45 -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] print information about key reuse location

2024-03-21 13:34:26 -07:00

profiler.py

[XLA:GPU] Add option to return FDO profile as textproto.

2024-03-26 01:35:27 -07: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

Evaluate the correctness of JAX complex functions using mpmath as a reference

2024-03-21 23:35:29 +02:00

tpu_custom_call.py

[Mosaic] Enable the lightweight stable serialization format

2024-03-21 02:52:40 -07:00

traceback_util.py

Upgrade remaining sources to Python 3.9

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

tree_util.py

[XLA:Python] Add a C++ implementation of flatten_one_level.

2024-03-21 15:57:23 -07: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

Removed obsolete call to libtpu_module.configure_library_path().

2024-03-26 16:06:41 -07:00