rocm_jax

clusters

Migrate JAX internals to builtin Python logging

2022-10-13 21:32:44 +02: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

Add examples to lax.psum to illustrate axis_index_groups better.

2022-12-23 12:04:59 -08:00

lib

Bump the minimum jaxlib version to 0.4.1.

2022-12-19 17:49:24 +00:00

nn

Remove typing_extensions dependency

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

numpy

[jax2tf] Improved error checking for jnp.squeeze with shape polymorphism

2022-12-20 17:58:32 +02:00

ops

Remove typing_extensions dependency

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

scipy

Added scipy.stats.bernoulli cdf and ppf.

2022-12-22 18:12:25 +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

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

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

ad_checkpoint.py

Merge pull request #13735 from jakevdp:private-linear-util

2022-12-21 05:15:38 -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

Move functions into api_util.py and dispatch.py to remove circular import error when pjit is imported in api.py for merging the jit and pjit frontend API.

2022-12-22 08:42:05 -08:00

api.py

Remove dependency of maps from pjit to avoid circular imports when importing pjit in api.py.

2022-12-22 13:35:23 -08:00

array.py

Remove more checks now that the minimum jaxlib version corresponds to xla_extension_version == 109. Also remove usage of xc._version and replace it with xla_extension_version.

2022-12-19 13:15:07 -08:00

basearray.py

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

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

basearray.pyi

[typing] fix Shard/Sharding type TODO

2022-12-07 09:44:03 -08:00

callback.py

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

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

checkify.py

Move jax.linear_util to jax._src.linear_util

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

cloud_tpu_init.py

Set JAX_PLATFORMS=tpu,cpu on TPUs

2022-10-05 04:39:04 +00:00

config.py

Revert #13747 since config values already have paired env vars.

2022-12-21 09:06:08 -08:00

core.py

Move jax.linear_util to jax._src.linear_util

2022-12-20 14:49:27 -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

Move functions into api_util.py and dispatch.py to remove circular import error when pjit is imported in api.py for merging the jit and pjit frontend API.

2022-12-22 08:42:05 -08:00

distributed.py

Migrate JAX internals to builtin Python logging

2022-10-13 21:32:44 +02:00

dlpack.py

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

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

dtypes.py

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

2022-12-16 21:00:14 -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

Remove dependency of maps from pjit to avoid circular imports when importing pjit in api.py.

2022-12-22 13:35:23 -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

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

2022-12-16 21:00:14 -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

[x64] make chack_grads() more type-safe

2022-12-02 12:51:41 -08:00

random.py

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

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

sharding.py

Remove more checks now that the minimum jaxlib version corresponds to xla_extension_version == 109. Also remove usage of xc._version and replace it with xla_extension_version.

2022-12-19 13:15:07 -08:00

source_info_util.py

Let is_user_frame ignore frames from stdlib.

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

stages.py

Bump the minimum jaxlib version to 0.4.1.

2022-12-19 17:49:24 +00:00

test_util.py

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

2022-12-16 21:00:14 -08: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

Merge pull request #13253 from NeilGirdhar:keypaths

2022-11-28 13:10:20 -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