rocm_jax

clusters

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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

[export] Refactor the imports for the public API of jax.experimental.export

2024-01-08 05:29:56 -08:00

interpreters

Remove indices and devices from shard_arg_handlers and shard_args.

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

lax

Remove reference to np.ComplexWarning

2024-01-08 09:53:06 -08:00

lib

[Mosaic] Remove Python implementation of apply_vector_layout and infer_memref_layout.

2024-01-10 13:00:21 -08:00

nn

Upgrade remaining sources to Python 3.9

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

numpy

[array API] add jnp.linalg.diagonal

2024-01-11 12:52:15 -08:00

ops

[shape_poly] Update the jax.ops.segment{max|...} to with with shape polymorphism

2023-12-19 12:02:39 +02:00

pallas

Fix all_gather_test

2024-01-11 14:38:43 -08:00

scipy

stats.binom.pmf: return zero for k > n

2024-01-02 10:53:44 -08:00

state

CI: bump mypy to version 1.8.0

2024-01-10 10:20:55 -08:00

third_party

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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

del add_any_p and zeros_like_p, replace aval-dispatched traceable

2023-12-21 17:04:21 -08:00

ad_checkpoint.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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

Always flatten args and kwargs together i.e. tree_flatten((args, kwargs)) so that we have a uniform in_tree structure everywhere.

2023-12-19 17:32:07 -08:00

api.py

Remove indices and devices from shard_arg_handlers and shard_args.

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

array.py

Remove indices and devices from shard_arg_handlers and shard_args.

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

basearray.py

Upgrade remaining sources to Python 3.9

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

basearray.pyi

Ensure sharding-related array properties are documented

2023-11-03 09:56:33 -07:00

cache_key.py

Remove the old cache-key generation code.

2023-12-12 16:34:32 -08:00

callback.py

Add a registry for primitives that require device_assignment during lowering

2023-12-08 16:31:41 -08:00

checkify.py

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

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

cloud_tpu_init.py

Remove calling configure_library_path during jax import and get libtpu path from libtpu_module.get_library_path().

2023-10-10 10:59:37 -07:00

compilation_cache_interface.py

Move implementation of compilation cache out of jax/experimental and into jax/_src.

2023-04-17 08:35:53 -07:00

compilation_cache.py

Correct the cache miss metric instrumentation due to the new min cache entry size flag

2024-01-08 14:03:33 -08:00

compiler.py

[pjrt:cpu] Add CpuTopology to TfrtCpuClient and enable persistent compilation cache for cpu backend

2024-01-10 12:40:57 -08:00

config.py

Introduce min entry size check for compilation cache.

2024-01-04 15:17:05 -08:00

core.py

[shape_poly] Replace non_negative_dim with max_dim and min_dim.

2024-01-08 20:54:18 +02: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

Upgrade remaining sources to Python 3.9

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

custom_derivatives.py

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

2023-12-21 22:15:52 -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

Make deprecation warnings warn at appropriate stacklevel

2023-05-23 14:43:38 -07:00

dispatch.py

Remove indices and devices from shard_arg_handlers and shard_args.

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

distributed.py

Upgrade remaining sources to Python 3.9

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

dlpack.py

DOC: add warning about dlpack and buffer mutation

2024-01-03 13:31:57 -08:00

dtypes.py

Support float8 in reduce_min & reduce_max

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

effects.py

Effects is now an immutable set

2023-12-11 08:45:52 -08:00

environment_info.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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

jaxpr_util.py

Upgrade remaining sources to Python 3.9

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

lax_reference.py

add an exp2 primitive and lax.exp2

2023-07-28 12:33:49 -07: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

add jax.explain_cache_misses tracing cache miss explanations

2023-12-26 21:54:27 -08:00

logging_config.py

Add jax_debug_log_modules config option.

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

maps.py

Delete mesh_sharding_specs from JAX

2024-01-02 11:14:39 -08:00

mesh.py

Upgrade most .py sources to 3.9

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

monitoring.py

Upgrade remaining sources to Python 3.9

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

op_shardings.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00: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 indices and devices from shard_arg_handlers and shard_args.

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

pretty_printer.py

Upgrade remaining sources to Python 3.9

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

prng.py

Remove indices and devices from shard_arg_handlers and shard_args.

2024-01-05 14:17:14 -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

Fix duplicate word occurrences

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

shard_alike.py

Rename batch_dims variable in shard_alike

2024-01-03 20:59:39 -08:00

sharding_impls.py

Remove indices and devices from shard_arg_handlers and shard_args.

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

sharding_specs.py

Delete sharding spec to HloSharding conversion since it's not used anymore.

2024-01-02 13:13:23 -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

Upgrade remaining sources to Python 3.9

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

stages.py

Fix indexing bug when querying _input_layouts

2023-11-27 21:12:07 -08:00

test_util.py

Add verbose kwarg to assertArraysEqual

2024-01-11 14:11:53 -08:00

tpu_custom_call.py

[Mosaic] Remove Python implementation of apply_vector_layout and infer_memref_layout.

2024-01-10 13:00:21 -08:00

traceback_util.py

Upgrade remaining sources to Python 3.9

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

tree_util.py

Remove a number of deprecated APIs

2023-12-18 10:08:47 -08:00

typing.py

[shape_poly] Simplify the API for processing polymorphic_shape specifications

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

util.py

Upgrade remaining sources to Python 3.9

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

xla_bridge.py

Guard call to os.register_at_fork() to unbreak Windows CI.

2024-01-04 08:10:32 -08:00