rocm_jax

clusters

Fix typos across the package

2023-09-22 14:54:31 -07:00

debugger

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

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

extend

move wrap_key_data to jax.random

2023-09-18 14:38:22 -07:00

image

Apply pyupgrade --py39-plus.

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

internal_test_util

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

interpreters

[jax_export] Add the next part of multi-platform lowering support.

2023-10-02 10:57:35 +03:00

lax

MAINT Reorder the overloads for lax.sort

2023-09-28 12:51:36 +01:00

lib

[CUDA] Add code to jax initialization that verifies that the CUDA libraries that are found are at least as new as the versions against which JAX was built.

2023-09-27 11:28:40 -07:00

nn

Update internal callers to avoid PRNGKeyArray

2023-09-13 14:05:42 -07:00

numpy

jnp.mean: for f16 inputs, accumulate in f32

2023-09-27 18:27:19 -07:00

ops

Internal: add dtypes.safe_to_cast utility & use to generate indexing warning

2023-09-07 12:18:14 -07:00

pallas

Correct typo in dtype ValueError() call.

2023-09-29 10:31:58 -07:00

scipy

Merge pull request #17733 from superbobry:dict-literals

2023-09-26 09:22:37 -07:00

state

Fix typos across the package

2023-09-22 14:54:31 -07:00

third_party

Add a type stub for jax.numpy.

2023-08-22 11:50:49 -07: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

Apply pyupgrade --py39-plus.

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

ad_checkpoint.py

[callbacks] Add support for shardable ordered effects.

2023-09-18 02:50:25 -07:00

ad_util.py

Apply pyupgrade --py39-plus.

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

api_util.py

Fix resolve_argnums if inspect.signature fails. In this case, donate_argnames was None leading an error in assert_no_intersection.

2023-07-31 12:10:15 -07:00

api.py

Introduce a LoweringParameters dataclass for easier plumbing

2023-09-29 08:23:05 +03:00

array.py

Better errors for array scalar/boolean conversion

2023-09-19 09:00:19 -07:00

basearray.py

Split basearray into separate Bazel module.

2023-03-13 11:14:41 -07:00

basearray.pyi

[typing] fix a few array type declarations

2023-09-12 13:21:48 -07:00

cache_key.py

Fix typos across the package

2023-09-22 14:54:31 -07:00

callback.py

[callbacks] Add support for shardable ordered effects.

2023-09-18 02:50:25 -07:00

checkify.py

custom prng: introduce mechanism to identify key arrays by dtype

2023-07-21 12:27:32 -07:00

cloud_tpu_init.py

Default to PJRT TPU runtime instead of StreamExecutor on older jaxlibs.

2023-08-23 00:06:16 +00: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

Set up plumbing for adding new compilation-cache-key generation algorithm.

2023-08-09 18:16:22 -07:00

compiler.py

Modify the name for JAX compilation cache adoption metric.

2023-09-13 16:44:10 -07:00

config.py

Reverts 1a9c94e6265b40c8f1de1e6d920208d648d70fdd

2023-09-27 06:58:35 -07:00

core.py

Improve behavior of core.valid_jaxtype

2023-09-22 13:46:09 -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

Fix type failures under an upcoming pytype change.

2023-04-07 07:10:09 -07:00

custom_derivatives.py

Fix typos across the package

2023-09-22 14:54:31 -07:00

custom_transpose.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

debugging.py

Apply pyupgrade --py39-plus.

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

deprecations.py

Make deprecation warnings warn at appropriate stacklevel

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

dispatch.py

Add temporary flag for forcing arg tuplization of lowered functions.

2023-09-30 23:14:38 -07:00

distributed.py

Fix typos across the package

2023-09-22 14:54:31 -07:00

dlpack.py

Fix dlpack type signatures to match Array API spec.

2023-09-08 10:12:32 -04:00

dtypes.py

jnp.issubdtype: refactor for better performance

2023-09-26 11:26:28 -07:00

effects.py

[callbacks] Add support for shardable ordered effects.

2023-09-18 02:50:25 -07:00

environment_info.py

Split _src files custom_api_util, deprecations, effects and environment_info into separate Bazel targets.

2023-03-10 12:26:05 -08:00

errors.py

Fix typos across the package

2023-09-22 14:54:31 -07:00

flatten_util.py

Avoid imports from the public jax.* namespace in more places internally.

2023-04-04 11:41:40 -07: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

Move jax.jaxpr_util to jax._src.jaxpr_util, and split it into a separate build target.

2023-08-08 10:09:09 -07:00

lax_reference.py

add an exp2 primitive and lax.exp2

2023-07-28 12:33:49 -07:00

lazy_loader.py

Apply pyupgrade --py39-plus.

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

linear_util.py

JEX: move jax.linear_util to jax.extend.linear_util

2023-08-30 18:32:12 -07:00

logging_config.py

Add jax_debug_log_modules config option.

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

maps.py

Introduce a LoweringParameters dataclass for easier plumbing

2023-09-29 08:23:05 +03:00

mesh.py

Fix mypy error

2023-08-29 13:25:12 -07:00

monitoring.py

Fix typing for kwargs.

2023-09-28 15:03:20 -07:00

op_shardings.py

Apply pyupgrade --py39-plus.

2023-07-21 14:49:44 -04: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

Introduce a LoweringParameters dataclass for easier plumbing

2023-09-29 08:23:05 +03:00

pretty_printer.py

Update flags to use the ABSL typed flag API.

2023-07-27 12:15:58 -07:00

prng.py

[random] add itemsize property to custom PRNG

2023-09-25 08:52:26 -07:00

profiler.py

Fix typos across the package

2023-09-22 14:54:31 -07:00

public_test_util.py

Move --jax_test_dut into the non-public test utilities.

2023-09-06 13:21:48 -07:00

random.py

[random] deprecate named key creation functions

2023-09-21 13:57:49 -07:00

sharding_impls.py

Make is_fully_addressable an abstract method and implement it on each concrete Sharding.

2023-09-05 17:28:22 -07:00

sharding_specs.py

Replace references to DeviceArray with Array.

2023-08-18 17:46:00 -04:00

sharding.py

Make is_fully_addressable an abstract method and implement it on each concrete Sharding.

2023-09-05 17:28:22 -07:00

source_info_util.py

Apply pyupgrade --py39-plus.

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

stages.py

[JAX] Fix a regression in cost_analysis API access for an alternative JAX backend

2023-09-29 21:42:33 -07:00

test_util.py

Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test().

2023-09-27 12:10:43 -07:00

tpu_custom_call.py

[Mosaic] Add support for specifying estimated costs for Mosaic kernels

2023-09-14 02:50:44 -07:00

traceback_util.py

Tracebacks no longer have JAX-internal frames prepended by default

2023-08-03 11:38:38 -07:00

tree_util.py

fast dispatch for functions over typed PRNG key arrays

2023-09-13 09:43:58 -07:00

typing.py

Copybara import of the project:

2023-07-24 14:38:20 -07:00

util.py

[typing] better annotations for jnp.ufunc

2023-09-21 12:48:14 -07:00

xla_bridge.py

[XLA: Python] Update PJRT plugin configuration mapping with bool type

2023-10-02 10:08:34 -07:00