Yash Katariya 1cb8d31c66 Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays. ...

Also comment out key reuse check in cpp dispatch since it's True for jax tests which prevent prng keys from taking Cpp dispatch.

PiperOrigin-RevId: 613289252

2024-03-06 11:42:40 -08:00

..

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

Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays.

2024-03-06 11:42:40 -08:00

lax

Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays.

2024-03-06 11:42:40 -08:00

lib

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

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

nn

Prevent accidental upcasting in jax.nn.initializers.

2024-03-01 14:24:26 -08:00

numpy

jnp.array: fix failure under numpy 2.0 copy semantics

2024-03-04 10:39:38 -08:00

ops

jax.numpy: implement scalar boolean indexing

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

pallas

[Pallas] Pass in compiler params via explicit compiler_params argument instead of passing via **kwargs

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

scipy

DOC: fix two minor doc issues

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

state

[mutable-arrays] move MutableArray, add eager, improve tests, fix bug

2024-03-01 15:03:23 -08: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

Axis names are now tracked via an effect

2024-03-04 05:42:03 -08:00

array.py

Add copy argument to Array.__array__

2024-03-05 09:31:16 -08: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

Fix bug where axis size selected incorrectly in pure_callback vmap rule

2024-02-26 15:09:20 -08: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

To avoid the inconsistency between process_index and process_id, replace backend.process_index with distributed.global_state.process_id in Jax compilation _cache_write function.

2024-02-15 10:48:05 -08:00

config.py

add upgrade (aka to-be-removed) flag for new select rule

2024-03-01 09:50:01 -08:00

core.py

Axis names are now tracked via an effect

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

[xmap-removal] remove reduce_axes from grad / vjp / backward_pass

2024-02-25 15:50:54 -08: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 typos across the package

2023-09-22 14:54:31 -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

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

Axis names are now tracked via an effect

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

Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays.

2024-03-06 11:42:40 -08:00

pretty_printer.py

Replaced most usages of abc.ABC with util.StrictABC

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

prng.py

Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays.

2024-03-06 11:42:40 -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

remove dead code in random

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

shard_alike.py

Guard shard_alike usage on xla_extension_version

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

sharding_impls.py

Make sure that if gspmd_sharding1 == gspmd_sharding2, then their hash also is equal.

2024-03-05 16:36:49 -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

Improve pytree mismatch error in AOT

2024-03-04 13:15:32 -08:00

test_util.py

Convert in_shardings to physical shardings in cpp dispatch path because the same happens with prng arrays.

2024-03-06 11:42:40 -08:00

tpu_custom_call.py

Import submodules from jax._src explicitly, instead of relying on import side-effects. It will lead to the missing x-refs in code search according to go/pywald-sawmill-analysis.

2024-02-06 15:47:16 -08: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

Small optimization to trace_context.

2024-02-12 14:30:11 -08:00

xla_bridge.py

Update

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