George Necula 4e89d43a75 Added JAX pytrees notebook ...

Also added docstrings to the tree_util module.

2019-11-24 20:29:07 +01:00

..

experimental

Implemented suggestions from code review.

2019-11-18 11:39:58 +01:00

interpreters

Remove as_abstract_val.

2019-11-22 10:52:38 -08:00

lax

Implement batching rule for lax._select_and_gather_add (#1736)

2019-11-21 11:52:58 -05:00

lib

Add jax.devices(), etc. to the docs.

2019-11-22 11:03:42 -08:00

nn

fix orthogonal initializer and reenable test

2019-10-29 11:34:20 -07:00

numpy

Add jax.numpy.dtype as an alias of numpy.dtype. (#1750)

2019-11-22 16:06:56 -05:00

ops

Temporarily remove jit decorator on gather/scatter ops.

2019-09-16 13:57:07 -07:00

scipy

Add jit decorators to most functions in jax.scipy.linalg. (#1741)

2019-11-21 16:31:16 -05:00

tools

Remove licenses() rule comment in BUILD files.

2019-09-26 14:54:07 -07:00

__init__.py

Import nn in jax/__init__.py.

2019-11-01 09:28:48 +00:00

abstract_arrays.py

Add bfloat16 support to JAX. (#1720)

2019-11-20 22:43:46 -05:00

ad_util.py

De-tuplify the rest of the core

2019-08-21 13:21:20 -07:00

api_util.py

make vmap structured axes work for any pytree

2019-10-31 14:09:12 -07:00

api.py

Move internal type-related functions into a new (internal) jax.types … (#1695)

2019-11-15 10:02:51 -05:00

BUILD

Add build target for optix

2019-11-01 15:39:19 +00:00

config.py

fix config.update when absl is enabled

2019-10-02 16:08:05 -07:00

core.py

fix reference cycle in jaxpr tracing using weakrefs

2019-11-19 15:23:08 -08:00

dtypes.py

Add bfloat16 support to JAX. (#1720)

2019-11-20 22:43:46 -05:00

flatten_util.py

prevent jit from treating keyword args as static

2019-04-10 22:09:14 -07:00

lax_linalg.py

Move internal type-related functions into a new (internal) jax.types … (#1695)

2019-11-15 10:02:51 -05:00

lax_reference.py

Add bfloat16 support to JAX. (#1720)

2019-11-20 22:43:46 -05:00

linear_util.py

improve names, avoid double lookup (thanks @hawkinsp)

2019-10-31 16:26:29 -07:00

pprint_util.py

source sync

2018-11-21 20:22:54 -08:00

random.py

replace x.shape with onp.shape(x) in random.py

2019-11-22 10:59:31 -08:00

test_util.py

Add bfloat16 support to JAX. (#1720)

2019-11-20 22:43:46 -05:00

tree_util.py

Added JAX pytrees notebook

2019-11-24 20:29:07 +01:00

util.py

Rm duplicates from end_nodes in toposort

2019-10-01 17:56:44 +01:00

version.py

bump version for pypi

2019-11-19 07:02:29 -08:00