Yash Katariya 6072d5993e Any devices passed to jax.sharding.Mesh are required to be hashable. ...

This is true for mock devices or user specific devices and jax.devices() too.

Fix the tests so that the mock devices are hashable.

PiperOrigin-RevId: 561103167

2023-08-29 12:20:54 -07:00

..

_src

Any devices passed to jax.sharding.Mesh are required to be hashable.

2023-08-29 12:20:54 -07:00

example_libraries

Use lower-case PEP 585 names for types.

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

experimental

Make mlir.custom_call() more general and expose it as jax.interpreters.mlir.custom_call().

2023-08-29 08:50:07 -07:00

extend

add jax.extend.random.wrap_key_data

2023-08-26 11:39:25 -07:00

image

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

interpreters

Make mlir.custom_call() more general and expose it as jax.interpreters.mlir.custom_call().

2023-08-29 08:50:07 -07:00

lax

remove deprecated jax.lax.prod function

2023-08-24 10:13:59 -07:00

lib

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

nn

Add test framework for module attribute

2023-04-21 13:20:16 -07:00

numpy

Give jax.numpy.array the type Callable.

2023-08-28 10:41:07 -07:00

ops

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

scipy

Deprecate jax.numpy.trapz.

2023-08-25 09:04:13 -06:00

tools

Load Python rules from rules_python

2023-08-24 10:22:57 -07:00

__init__.py

Update __init__.py

2023-08-28 10:13:16 +05:00

abstract_arrays.py

Mark jax.abstract_arrays as deprecated

2023-06-07 23:36:40 -07:00

ad_checkpoint.py

…

api_util.py

…

BUILD

Merge pull request #17307 from froystig:wrap-key

2023-08-27 12:58:50 -07:00

cloud_tpu_init.py

…

collect_profile.py

…

config.py

…

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_batching.py

…

custom_derivatives.py

custom_vjp symbolic zeros support, take two

2023-04-05 11:17:05 -07:00

custom_transpose.py

…

debug.py

Explicitly re-export all names from jax.debug

2023-04-07 07:24:59 -07:00

distributed.py

…

dlpack.py

…

dtypes.py

Copybara import of the project:

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

errors.py

Tracebacks no longer have JAX-internal frames prepended by default

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

flatten_util.py

ravel_pytree now accepted by static linters

2023-03-07 23:35:48 +00:00

linear_util.py

jax.linear_util: remove unnecessary exported names

2022-12-21 18:52:01 -08:00

monitoring.py

Move contents of jax.monitoring to jax._src.monitoring.

2023-02-08 15:03:22 -08:00

prng.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

profiler.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

py.typed

…

random.py

add random.triangular and random.lognormal

2023-08-05 12:27:53 +08:00

sharding.py

Delete OpShardingSharding export since it has been 3 months since it was deprecated. Also remove deprecation warnings for MeshPspecSharding.

2023-06-08 13:53:38 -07:00

stages.py

add CompilerOptions type alias for AOT compiler options

2023-04-12 16:49:13 -07:00

test_util.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

tree_util.py

Add register_static decorator to tree_util to facilitate creating leafless classes.

2023-08-21 16:54:29 -07:00

typing.py

jax.typing: recommend instance check in Python 3.10 or newer

2023-03-27 10:01:28 -07:00

util.py

Document why 'import name as name' is used

2022-12-14 15:07:04 -08:00

version.py

Add jax release jobs. Also make the __version__ be the .dev version by default unless JAX_RELEASE env var is set.

2023-08-18 08:33:06 -07:00