Peter Hawkins 5527966b27 [JAX] Deprecate .to_py() property on arrays. Implement __array__ instead. ...

.to_py() was something of an accidental export from the JAX array classes. There are other mechanisms to turn a JAX array into a NumPy array, including `np.asarray(x)` and `jax.device_get(x)`. Deprecate this mechanism because it is redundant.

PiperOrigin-RevId: 469984029

2022-08-25 07:28:27 -07:00

..

debugger

Internal change

2022-08-18 11:39:07 -07:00

image

[x64] make image_test compatible with strict dtype promotion

2022-06-15 11:41:23 -07:00

lax

Upgrade logistic (sigmoid) function into a lax primitive.

2022-08-24 15:39:37 -07:00

lib

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00

nn

Upgrade logistic (sigmoid) function into a lax primitive.

2022-08-24 15:39:37 -07:00

numpy

Remove deprecated array.tile() method

2022-08-23 07:36:59 -07:00

ops

[x64] deprecate unsafe type casting in scatter-update operations

2022-06-09 15:21:49 -07:00

scipy

Upgrade logistic (sigmoid) function into a lax primitive.

2022-08-24 15:39:37 -07:00

state

Merge pull request #11900 from sharadmv:for-loop

2022-08-18 20:14:17 -07:00

third_party

[x64] make jax.scipy.signal compatible with strict dtype promotion.

2022-06-21 09:28:46 -07:00

__init__.py

Move most of the implementation of jax.scipy into jax._src.scipy.

2020-10-16 17:04:25 -04:00

abstract_arrays.py

[remove-units] prevent cond partial eval from introducing units

2022-04-24 14:28:33 -07:00

ad_checkpoint.py

in ad_checkpoint WrapHashably.__eq__, check _both_ are hashable

2022-08-23 19:06:57 -07:00

ad_util.py

Add stateful for loop primitives (#10982)

2022-06-15 15:55:38 -07:00

api_util.py

Fix common typo: Tuple[X] -> Tuple[X, ...]

2022-08-16 11:47:22 -04:00

api.py

Enable fast path in the Array constructor. This means that the rearranging of _arrays according to the device_assignment won't happen when fastpath is enabled because we assume that jax transformations will return the right arrangement.

2022-08-23 10:20:26 -07:00

callback.py

Apply suggestions from code review

2022-08-18 10:50:50 -07:00

checkify.py

Make checkify tests pass with Array and add methods on Array that are present on DA.

2022-08-16 16:52:06 -07:00

cloud_tpu_init.py

Internal change

2022-08-18 11:39:07 -07:00

config.py

Add upgrade=True to jax_array flag so that its marked as transient flag which will eventually be set to True.

2022-08-22 13:07:53 -07:00

custom_api_util.py

enable custom transformation "stacking"

2022-01-11 17:55:08 -08:00

custom_batching.py

Add spmd_axis_name to vmap to allow constraining mapped PartitionSpecs.

2022-08-08 19:41:42 -07:00

custom_derivatives.py

rolling forward deletion of custom_jvp_call_jaxpr_p yet again...

2022-08-18 14:02:40 -07:00

custom_transpose.py

[dynamic-shapes] initial support for dynamic shape typechecks

2022-06-17 14:57:19 -07:00

debugging.py

Allow debug prints in staged out custom derivative functions

2022-08-12 19:49:09 -07:00

device_array.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

dispatch.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

distributed.py

Merge pull request #11946 from nvcastet:fix_distributed_timeout

2022-08-24 07:28:44 -07:00

dlpack.py

Avoid import-time dependency on jax.experimental

2022-08-19 11:30:25 -07:00

dtypes.py

Merge pull request #11857 from jakevdp:fix-bool-args

2022-08-12 11:45:35 -07:00

errors.py

prepare to switch to new remat

2022-08-04 12:25:03 -07:00

flatten_util.py

Documentation improvements.

2022-03-08 09:37:33 -05:00

iree.py

[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.

2022-08-25 07:28:27 -07:00

lax_reference.py

Upgrade logistic (sigmoid) function into a lax primitive.

2022-08-24 15:39:37 -07:00

pretty_printer.py

Internal change

2022-08-18 11:39:07 -07:00

prng.py

support jax.experimental.array.Array as a base array for key arrays

2022-08-24 19:49:02 -07:00

profiler.py

jax.profiler: remove deprecated functions

2022-07-19 08:13:44 -07:00

public_test_util.py

Separate jax.test_util implementations into public and private sources.

2022-04-04 14:43:39 -07:00

random.py

roll-forward #11952, take 2

2022-08-22 13:57:31 -07:00

source_info_util.py

Only wrap one level of the name-stack with transform names.

2022-06-10 03:50:26 -07:00

stages.py

Fix common typo: Tuple[X] -> Tuple[X, ...]

2022-08-16 11:47:22 -04:00

test_util.py

roll-forward #11952, take 2

2022-08-22 13:57:31 -07:00

traceback_util.py

Fix Python 3.11 compatibility problems.

2022-08-10 19:45:24 +00:00

tree_util.py

promise to flatten trees in left-to-right order

2022-07-28 19:28:20 -07:00

util.py

Use c++ weakref LRU cache implementation as a drop in replacement for jax's

2022-08-18 14:36:08 -07:00