rocm_jax

clusters

Don't set TPU topology env vars in cloud_tpu_init.py

2023-01-13 22:02:16 +00:00

debugger

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

image

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lax

Raise an error when attempting to mutate Jaxpr objects

2023-01-23 09:37:58 -08:00

lib

CI: fix mypy jaxlib version

2023-01-24 06:57:23 -08:00

nn

Added shape error checking for compute_fans

2023-01-18 20:59:11 -08:00

numpy

[shape_poly] Generalize binary operations with symbolic dimensions.

2023-01-21 04:26:59 -08:00

ops

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

scipy

Adding info to CG and BICGSTAB

2023-01-22 21:47:34 +00:00

state

Add support for Ellipsis as an index for stateful operations.

2022-12-23 22:46:50 -08:00

third_party

Address drastic slowdown in mypy runtime

2022-12-16 14:48:26 -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

Expose fp8 in jax dtypes and mlir builder.

2023-01-13 18:12:12 -08:00

ad_checkpoint.py

tweak remat named policies

2023-01-04 08:25:57 -08:00

ad_util.py

Merge pull request #13693 from jakevdp:typing-ad-util

2022-12-20 16:50:11 -08:00

api_util.py

Make error_test a jax_test so that we can test other configs and fix it with jit/pjit merge.

2023-01-17 18:43:05 -08:00

api.py

Clarify the docstring for vjp.

2023-01-20 11:25:23 -05:00

array.py

Make Shard.device and Shard.data read-only properties.

2023-01-05 14:27:17 +00:00

basearray.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

basearray.pyi

[typing] add missing attributes to jax.Array

2023-01-11 10:00:59 -08:00

callback.py

Add IO callback

2023-01-17 13:55:05 -08:00

checkify.py

Make jit a thin wrapper around pjit which ignores the mesh context manager (just like how it is today)

2023-01-12 17:24:32 -08:00

cloud_tpu_init.py

Don't set TPU topology env vars in cloud_tpu_init.py

2023-01-13 22:02:16 +00:00

config.py

Add jax_jit_pjit_api_merge config to help the transition to merge jit and pjit.

2023-01-03 12:59:46 -08:00

core.py

Use __slots__ in Jaxpr & ClosedJaxpr

2023-01-23 11:33:30 -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

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

custom_derivatives.py

Minor docstring fix for custom_gradient.

2022-12-23 12:12:48 -08:00

custom_transpose.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

debugging.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

device_array.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

dispatch.py

Fix typo "invalud" -> "invalid" in error message.

2023-01-20 08:48:24 -08:00

distributed.py

Add Open MPI automatic distributed initialization

2023-01-11 17:08:09 -06:00

dlpack.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

dtypes.py

Expose fp8 in jax dtypes and mlir builder.

2023-01-13 18:12:12 -08:00

environment_info.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

errors.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

flatten_util.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

iree.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lax_reference.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

lazy_loader.py

Avoid loading scipy eagerly.

2022-10-12 19:51:09 +00:00

linear_util.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -08:00

path.py

Make the etils[epath] dependency of JAX optional.

2022-10-12 21:16:48 +00:00

pjit.py

Cache the creation of ClosedJaxpr in pjit_transpose which if not cached breaks the compilation cache.

2023-01-24 09:58:46 -08:00

pretty_printer.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

prng.py

Simplify Array's shard_arg_handler by merging pmap and pjit/xmap paths

2022-12-27 10:16:44 -08:00

profiler.py

Add a number of missing function cross-references in the docs.

2022-11-07 12:00:26 -05:00

public_test_util.py

Expose fp8 in jax dtypes and mlir builder.

2023-01-13 18:12:12 -08:00

random.py

migrate internal dependencies from jax.core to jax._src.core

2022-12-16 21:00:14 -08:00

sharding.py

Introduce is_equivalent_to method on Sharding to check if 2 shardings mean the same thing.

2023-01-23 11:04:08 -08:00

source_info_util.py

Let is_user_frame ignore frames from stdlib.

2022-11-08 10:50:08 -08:00

stages.py

Migrate JAX from producing MHLO to producing StableHLO

2022-12-27 08:53:20 -08:00

test_util.py

Add pjrt_c_api_unimplemented pytest marker to skip unsupported tests.

2023-01-12 22:17:23 +00:00

traceback_util.py

Add a number of missing function cross-references in the docs.

2022-11-07 12:00:26 -05:00

tree_util.py

specialize tree prefix error message for list/tuple

2023-01-20 10:51:02 -08:00

typing.py

Remove typing_extensions dependency

2022-12-05 15:42:26 -08:00

util.py

Remove local_imports of sharding.py. Adding pxla local imports but then cleaning those up will be super easy since those will be the only ones left and restricted to sharding.py file only.

2022-11-29 16:42:03 -08:00