rocm_jax

clusters

Apply pyupgrade --py39-plus.

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

debugger

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

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

image

Apply pyupgrade --py39-plus.

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

internal_test_util

add an exp2 primitive and lax.exp2

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

interpreters

Only donate if memory kinds match. This shouldn't break any existing behavior. Also only warn about unused donations if there are unused donations.

2023-08-15 13:26:18 -07:00

lax

Remove uses of XLA translation rules.

2023-08-15 12:53:36 -07:00

lib

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

nn

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/17027 from jakevdp:dtypes-annotations a116a9c498a7b085f9b3fec93b37da12289f6e31

2023-08-08 20:38:44 +00:00

numpy

jnp.ufunc: minor cleanups & test fixes

2023-08-14 15:19:46 -07:00

ops

Apply pyupgrade --py39-plus.

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

pallas

[jax_triton] Add user-specified name field to serialized format.

2023-08-16 02:53:51 -07:00

scipy

Remove superfluous double-where in xlogy and xlog1py.

2023-07-28 07:11:24 -07:00

state

Apply pyupgrade --py39-plus.

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

third_party

Add some more type annotations to lax_numpy.py.

2023-08-11 08:07:24 -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

Apply pyupgrade --py39-plus.

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

Move some utilities out of dispatch.py next to their users, add more types.

2023-08-08 10:52:11 -07:00

array.py

Reverts 7012a05497faf4d33c967bee3cebc83588234e63

2023-08-11 10:30:13 -07:00

basearray.py

Split basearray into separate Bazel module.

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

basearray.pyi

Relax axis type annotations for reduction methods on jax.Array.

2023-08-11 14:11:37 -07:00

cache_key.py

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

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

callback.py

Apply pyupgrade --py39-plus.

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

Remove option to use StreamExecutor Cloud TPU client in JAX

2023-08-08 14:05:27 -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

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

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

compiler.py

Move compiler APIs out of dispatch.py and xla_bridge.py into a new jax._src.compiler module.

2023-08-15 06:39:46 -07:00

config.py

[jax2tf] Bump the default JAX serialization version to 7.

2023-08-11 22:49:31 -07:00

core.py

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/17027 from jakevdp:dtypes-annotations a116a9c498a7b085f9b3fec93b37da12289f6e31

2023-08-08 20:38:44 +00: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

Apply pyupgrade --py39-plus.

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

Move compiler APIs out of dispatch.py and xla_bridge.py into a new jax._src.compiler module.

2023-08-15 06:39:46 -07:00

distributed.py

[JAX] Remove the non-coordination service distributed service implementation from JAX.

2023-08-07 15:17:25 -07:00

dlpack.py

[JAX] Delete _DeviceArray and DeviceArray.

2023-03-29 15:07:14 -07:00

dtypes.py

Support (u)int4 for jax.dtypes.scalar_type_of()

2023-08-09 13:52:04 -07:00

effects.py

Apply pyupgrade --py39-plus.

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

errors: create TracerBoolConversionError for more targeted debugging tips

2023-06-21 01:41:45 -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

Apply pyupgrade --py39-plus.

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

logging_config.py

Add jax_debug_log_modules config option.

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

maps.py

Move functions out of xla.py closer to their users.

2023-08-08 14:40:42 -07:00

mesh.py

[JAX] Use DeviceList in JAX Sharding implementations

2023-08-14 18:11:23 -07:00

monitoring.py

Add functions to unregister event duration listeners.

2023-08-09 10:43:02 -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

Canonicalize to default memory in init of Shardings only on the backends that support memories right now.

2023-08-04 16:27:15 -07:00

pretty_printer.py

Update flags to use the ABSL typed flag API.

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

prng.py

de-emphasize internal array implementation type in key array repr

2023-08-14 12:48:19 -07:00

profiler.py

Split _src/profiler into a separate BUILD target.

2023-03-10 11:38:46 -08:00

public_test_util.py

Add support for float8_e4m3fnuz and float8_e5m2fnuz.

2023-08-07 11:48:53 +01:00

random.py

Merge pull request #16960 from JiaYaobo:add_random_lognormal_and_triang

2023-08-07 11:14:04 -07:00

sharding_impls.py

[JAX] Use DeviceList in JAX Sharding implementations

2023-08-14 18:11:23 -07:00

sharding_specs.py

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/17027 from jakevdp:dtypes-annotations a116a9c498a7b085f9b3fec93b37da12289f6e31

2023-08-08 20:38:44 +00:00

sharding.py

[Memories] Add Memories support to jax.jit and jax.device_put!

2023-08-04 09:44:24 -07:00

source_info_util.py

Apply pyupgrade --py39-plus.

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

stages.py

Apply pyupgrade --py39-plus.

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

test_util.py

Merge pull request #16965 from jakeh-gc:fp8_fnuz

2023-08-07 10:43:18 -07:00

tpu_custom_call.py

Add fields kernel_name and kernel_regeneration_metadata to tpu_custom_call backend config

2023-08-14 02:35:20 -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

Bump minimum_jaxlib_version to 0.4.14. xla_extension_version is 174 and mlir_api_version is 54

2023-08-01 08:53:28 -07:00

typing.py

Copybara import of the project:

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

util.py

ENH: Update numpy exceptions imports

2023-08-07 19:08:41 +02:00

xla_bridge.py

[PJRT C API] Let framework explicitly check whether a plugin is initialized and initialize the plugin.

2023-08-15 15:24:14 -07:00