Matthew Johnson 0a693faf48 add pjit forwarding rule ...

Co-authored-by: Roy Frostig <frostig@google.com>

2024-05-25 17:46:01 +00:00

..

clusters

Add option to set coordinator lookup timeout for TPU clusters

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

cudnn

CI: fix mypy error

2024-05-20 13:23:15 -07:00

debugger

Merge pull request #20858 from rajasekharporeddy:doc_typos

2024-04-25 06:58:27 -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] Add backwards compatibility test for Pallas call on GPUs.

2024-05-02 05:38:33 -07:00

interpreters

add pjit forwarding rule

2024-05-25 17:46:01 +00:00

lax

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

lib

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

nn

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

numpy

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

ops

Support scalar boolean indices in arr.at[idx].set(vals)

2024-05-20 05:33:36 -07:00

pallas

Added support for printing scalar values in Pallas TPU kernels

2024-05-23 10:02:00 -07:00

scipy

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

state

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

third_party

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01: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

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

ad_checkpoint.py

Initialize JaxprEqnContext only in new_jaxpr_eqn and new_eqn_recipe with the current active compute type if no ctx is specified.

2024-05-22 15:16:58 -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

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

api.py

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

array.py

Improve error message when trying to fetch value of non-addressable array.

2024-05-23 12:41:34 -07:00

basearray.py

Upgrade remaining sources to Python 3.9

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

basearray.pyi

Add __buffer__ and __releasebuffer__ to the Array type stub.

2024-05-22 14:17:14 +00: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

Clarified the type of the inputs to callback APIs

2024-05-16 11:29:09 -07:00

checkify.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

cloud_tpu_init.py

Change determination of cloud TPU to check for TPU chips.

2024-05-02 07:22:56 -07:00

compilation_cache_interface.py

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

2024-04-16 20:28:24 -07:00

compilation_cache.py

Increase some log levels to info in public methods.

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

compiler.py

Bumped mypy to 1.10.0 and ruff to 0.4.4

2024-05-16 23:16:32 +01:00

compute_on.py

Support remat + compute_on. If the rematted computation is annotated to run on host, the backward pass will also execute on host. Also enable no-op nested compute tests.

2024-05-17 18:59:49 -07:00

config.py

reintroduce the Threefry GPU kernel lowering, under a flag

2024-05-01 10:33:31 -07:00

core.py

Initialize JaxprEqnContext only in new_jaxpr_eqn and new_eqn_recipe with the current active compute type if no ctx is specified.

2024-05-22 15:16:58 -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

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

custom_derivatives.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

custom_transpose.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

debugging.py

Merge pull request #21261 from superbobry:mypy-ruff

2024-05-17 00:10:27 -07:00

deprecations.py

Refactor & test internal deprecation APIs

2024-05-20 17:16:31 -07:00

dispatch.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

distributed.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

dlpack.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

dtypes.py

Refactor array_api namespace, relying more directly on jax.numpy

2024-05-02 18:17:45 +00:00

earray.py

start adding EArray, a jax.Array analog that can contain extended dtypes

2024-04-06 13:09:25 -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

ffi.py

Package XLA FFI headers with jaxlib wheel

2024-05-22 12:28:38 -04:00

flatten_util.py

Fix duplicate word occurrences

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

gfile_cache.py

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

2024-04-16 20:28:24 -07: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

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

lax_reference.py

Implements Ragged Dot API

2024-05-11 06:40:18 -07:00

layout.py

Prefer raising of TypeError for invalid types instead of ValueError.

2024-04-08 13:08:24 +05:30

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

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07: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

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

partition_spec.py

Improve the sharding documentation.

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

path.py

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

2024-04-16 20:28:24 -07:00

pickle_util.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

pjit.py

add pjit forwarding rule

2024-05-25 17:46:01 +00:00

pretty_printer.py

Compute source maps when pretty-printing jaxprs.

2024-05-06 15:45:25 -04:00

prng.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

profiler.py

Fix doc Typos

2024-04-22 10:32:51 +05:30

public_test_util.py

Support float8 in reduce_min & reduce_max

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

random.py

Finalize deprecation of batched keys to PRNG functions

2024-05-22 09:40:32 -07:00

shard_alike.py

Bump minimum jaxlib version to 0.4.27

2024-05-07 16:07:59 -07:00

sharding_impls.py

Fix unhashable: list error when there was a list present inside PartitionSpec.

2024-05-22 13:56:19 -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

Remove indices and devices from shard_arg_handlers and shard_args.

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

source_info_util.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

sourcemap.py

Merge pull request #21096 from gspschmid:gschmid/sourcemaps

2024-05-08 05:44:08 -07:00

stages.py

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

test_util.py

Add jax.make_array_from_process_local_data to create a distributed tensor from host data and supporting scaffolding in sharding to be able to figure out dimensions of host data required.

2024-05-15 22:06:45 -07:00

tpu_custom_call.py

[Mosaic] Add guard for absl-py in tpu_custom_call.py.

2024-04-26 08:36:32 -07:00

traceback_util.py

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

tree_util.py

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00

tree.py

Add jax.tree module with aliases of jax.tree_util

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

typing.py

Add support for max_version, dl_device, copy kwargs in __dlpack__

2024-04-11 16:44:19 +00:00

util.py

Add discharge rules for scan with mutable arrays. Move mutable array tests to separate file.

2024-05-02 14:36:16 -04:00

xla_bridge.py

Merge pull request #21273 from superbobry:mypy-ruff

2024-05-22 06:35:38 -07:00