Skye Wanderman-Milne b68a08adf1 Add programmatic profiling APIs, and rename some existing APIs. ...

This change provides aliases for the renamed APIs so existing code
won't break. We should remove these aliases after the next release.

2021-03-25 11:50:55 -07:00

..

_src

Add programmatic profiling APIs, and rename some existing APIs.

2021-03-25 11:50:55 -07:00

experimental

Merge pull request #6203 from hawkinsp:x64

2021-03-24 12:58:20 -07:00

image

Move implementation of jax.image to jax._src.image.

2020-10-17 11:37:55 -04:00

interpreters

simplify cotangent dropping in ad.backward_pass

2021-03-24 16:25:50 -07:00

lax

add clz primitive

2021-03-19 22:54:36 -07:00

lib

Fix flake8 error.

2021-03-24 15:28:01 -04:00

nn

Delete jax.nn.functions.

2020-11-09 19:48:54 -05:00

numpy

Initial implementation of jnp.delete

2021-03-16 17:05:23 -07:00

ops

Move jax.ops implementation into jax._src.ops.

2020-10-17 11:45:28 -04:00

scipy

added bicgstab to new jax repo

2021-02-18 18:01:28 -08:00

tools

Merge pull request #5173 from jakevdp:conv-elem-type-cleanup

2020-12-11 15:27:40 -08:00

__init__.py

add jax_default_matmul_precision flag & context mngr

2021-03-24 14:03:58 -07:00

abstract_arrays.py

fix convert_element_type on large Py int inputs

2021-03-21 19:08:59 -07:00

ad_util.py

Some minor typing fixes

2021-03-02 09:42:52 -08:00

api_util.py

Move jax.tree_util implementation to jax._src.tree_util.

2021-03-24 12:00:38 -07:00

api.py

Switch disable_jit() to use the common boolean state mechanism.

2021-03-24 19:46:52 -04:00

BUILD

Merge pull request #5717 from google:dynamic-shapes2

2021-02-16 18:45:13 -08:00

cloud_tpu_init.py

Automatically initialize Cloud TPU topology env vars if running on a Cloud TPU VM.

2021-03-10 09:15:31 -08:00

config.py

Switch disable_jit() to use the common boolean state mechanism.

2021-03-24 19:46:52 -04:00

core.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

custom_derivatives.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

dlpack.py

[call_tf] Add support for DLPack to avoid copying arrays in eager mode

2021-02-13 12:49:51 +02:00

dtypes.py

jax.dtypes: Remove outdated bfloat16 checks

2021-03-24 13:14:53 -07:00

errors.py

errors: add NonConcreteBooleanIndexError & debugging tips

2021-03-23 11:23:20 -07:00

flatten_util.py

generalize ravel_pytree to handle int types, add tests

2021-03-19 10:50:02 -07:00

jax.bzl

Move jax.bzl into correct location. Unbreaks build.

2021-01-21 06:56:23 -08:00

jaxpr_util.py

[JAX] Move jax.util to jax._src_util.

2021-01-11 14:21:07 -08:00

lax_reference.py

add clz primitive

2021-03-19 22:54:36 -07:00

linear_util.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

profiler.py

Add programmatic profiling APIs, and rename some existing APIs.

2021-03-25 11:50:55 -07:00

py.typed

Improve type annotations

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

random.py

DOC: fix jax.random documentation

2020-11-20 16:41:32 -08:00

test_util.py

unify configuration state handling

2021-03-23 18:56:01 -07:00

tree_util.py

Move jax.tree_util implementation to jax._src.tree_util.

2021-03-24 12:00:38 -07:00

util.py

Export split_dict from jax.util

2021-01-29 11:13:36 -08:00

version.py

Increase minimum jaxlib version to 0.1.64.

2021-03-24 13:49:48 -04:00