Peter Hawkins 9bf476aaad Redefine jnp.DeviceArray as jax.Array. ...

The concrete DeviceArray class is slated for deletion on March 15 as part of the jax.Array migration. Replace jnp.DeviceArray with its superclass (jax.Array), which is the closest equivalent in the new world.

In the future, jnp.DeviceArray will be deprecated and deleted in favor of jax.Array.

PiperOrigin-RevId: 515432432

2023-03-09 13:56:13 -08:00

..

_src

Split _src modules cloud_tpu_init, lazy_loader, path, monitoring into their own pytype_library Bazel targets.

2023-03-09 13:11:04 -08:00

example_libraries

Link to JAXopt in example_libraries.optimizers

2023-01-25 10:54:31 -08:00

experimental

[jax2tf] Work around a bug in jax2tf tests, to unblock the test

2023-03-09 12:57:48 -08:00

image

Document why 'import name as name' is used

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

interpreters

add result info to mhlo, fixes #14780

2023-03-06 21:21:26 -08:00

lax

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

lib

[JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge.

2023-02-28 07:01:57 -08:00

nn

Fix __module__ of jax.nn.initializers.* to be jax.nn.initializers.

2023-02-22 14:13:15 -05:00

numpy

Redefine jnp.DeviceArray as jax.Array.

2023-03-09 13:56:13 -08:00

ops

Document why 'import name as name' is used

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

scipy

Roll back #14792

2023-03-07 18:31:19 -08:00

tools

Prepare for jax and jaxlib 0.4.4 release

2023-02-16 08:37:15 -08:00

__init__.py

Shorten alias chains for names exported in jax. namespace.

2023-03-01 09:19:44 -08: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

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

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

api_util.py

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

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

BUILD

Split _src modules cloud_tpu_init, lazy_loader, path, monitoring into their own pytype_library Bazel targets.

2023-03-09 13:11:04 -08:00

cloud_tpu_init.py

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

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

collect_profile.py

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

config.py

Cleanup: remove unused imports in private modules

2022-10-20 14:37:21 -07:00

core.py

Hide accidental exports from jax.core.

2023-02-21 17:48:40 -08:00

custom_batching.py

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

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

custom_derivatives.py

custom_jvp symbolic zeros support

2023-02-24 07:33:49 -08:00

custom_transpose.py

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

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

debug.py

Add inspect_array_sharding, enabling looking at shardings in pjit-ted functions

2022-09-22 17:36:56 -07:00

distributed.py

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

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

dlpack.py

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

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

dtypes.py

Remove jax.dtypes._jax_types

2023-02-01 12:49:06 -08:00

errors.py

Document why 'import name as name' is used

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

flatten_util.py

ravel_pytree now accepted by static linters

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

jaxpr_util.py

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

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

Improve type annotations

2021-01-13 10:26:35 -05:00

random.py

add random.chisquare and random.f

2023-03-01 15:03:50 +08:00

sharding.py

Rename jax.sharding.OpShardingSharding to jax.sharding.GSPMDSharding. jax.sharding.OpShardingSharding will be removed in 3 months from Feb 17, 2023.

2023-02-17 17:11:06 -08:00

stages.py

Expose ArgInfo

2023-03-08 10:03:17 -08:00

test_util.py

Document why 'import name as name' is used

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

tree_util.py

Internal change

2023-03-04 13:24:11 +00:00

typing.py

DOC: improve usage recommendation in jax.typing

2023-02-21 04:58:21 -08:00

util.py

Document why 'import name as name' is used

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

version.py

Bump version and changelog after jax 0.4.5 release

2023-03-02 16:08:34 -08:00