Peter Hawkins 74f1ab0503 Export Device as jax.Device. ...

Users are writing things like jax.lib.xla_client.Device in type annotations which is not a public API. Add a supported public name for the Device type.

2023-02-02 12:58:15 -05:00

..

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

interpreters

Reapply: move jax.interpreters.ad to jax._src.interpreters.ad

2023-02-02 09:29:05 -08:00

lax

Clean up: remove unused helper functions

2023-02-01 09:55: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

Clean up: remove unused helper functions

2023-02-01 09:55:58 -08:00

ops

Remove typing_extensions dependency

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

scipy

Revert "Merge pull request #14113 from botev:main"

2023-02-01 19:50:27 -08: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

Export Device as jax.Device.

2023-02-02 12:58:15 -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

Merge pull request #14250 from mattjj:checkify-retracing

2023-02-01 16:44:56 -08:00

cloud_tpu_init.py

Add environment variable check in additional to libtpu check for cloud tpu vm

2023-01-25 09:51:37 -08:00

config.py

Fix the warning being raised when jax.Array is True about using jax.Array

2023-01-27 10:20:44 -08:00

core.py

comment on opaque dtypes

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

Clean up: remove unused helper functions

2023-02-01 09:55:58 -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

Merge pull request #14248 from jakevdp:dead-code

2023-02-01 21:25:46 +00: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

Merge pull request #14248 from jakevdp:dead-code

2023-02-01 21:25:46 +00: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

Clean up: remove unused helper functions

2023-02-01 09:55:58 -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

Roll-back https://github.com/google/jax/pull/14144 due to downstream test failures

2023-01-25 12:15:36 -08:00

random.py

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

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

sharding.py

Begin deprecation of public jax.ShapedArray

2023-01-30 11:27:58 -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

Remove the usage of _arrays from tests

2023-01-30 20:02:37 -08:00

traceback_util.py

add blank line, mainly to trigger/test source sync

2023-02-01 13:56:29 -08: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 unused CacheInfo namedtuple

2023-01-31 11:36:43 -08:00