Matthew Johnson c41c4de5a5

lower fori_loop to scan when possible (#2414) ...

When a fori_loop specialized on trip count is to be evaluated, it's
preferable to generate a scan rather than a while_loop because the
former is reverse-mode differentiable while the latter is not. Otherwise
they're essentially the same; in particular, no extensive inputs/outputs
arise unless reverse-mode autodiff is applied.

Also fixes #2412.

2020-03-13 15:15:55 -07:00

..

experimental

Add a module to apply updates every k steps (and accumulate them otherwise) (#2350)

2020-03-10 06:40:38 -07:00

interpreters

remove input shapes from params of some primitives (#2410)

2020-03-13 07:13:29 -07:00

lax

lower fori_loop to scan when possible (#2414)

2020-03-13 15:15:55 -07:00

lib

Fix exception causes all over the codebase (#2376)

2020-03-09 16:06:12 -04:00

nn

Removed unused imports (#2385)

2020-03-09 20:42:08 +01:00

numpy

jit lax_numpy.roll (#2392)

2020-03-09 13:21:30 -07:00

ops

Removed unused imports (#2385)

2020-03-09 20:42:08 +01:00

scipy

Rename np -> jnp and onp -> np in jax/scipy to match the standard JAX… (#2370)

2020-03-07 10:43:51 -05:00

third_party

Fixed link to Google 3rd party OSS components

2020-02-10 19:05:11 +01:00

tools

Remove "from __future__" uses from JAX. (#2117)

2020-01-29 12:29:03 -05:00

__init__.py

Import nn in jax/__init__.py.

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

abstract_arrays.py

Give Vars an aval. (#2299)

2020-03-09 10:14:23 +01:00

ad_util.py

Removed unused imports (#2385)

2020-03-09 20:42:08 +01:00

api_util.py

Remove "from __future__" uses from JAX. (#2117)

2020-01-29 12:29:03 -05:00

api.py

allow vmap in_axes to be a list, fixes #2367 (#2395)

2020-03-10 08:29:46 -07:00

BUILD

Add missing sources to jax build. (#2208)

2020-02-10 11:40:36 -05:00

config.py

Removed dependency on distutils.strtobool

2020-02-06 17:27:46 +01:00

core.py

Fix exception causes all over the codebase (#2376)

2020-03-09 16:06:12 -04:00

dlpack.py

Initial implementation of DLPack support. (#2123)

2020-01-29 23:19:14 -05:00

dtypes.py

Remove "from __future__" uses from JAX. (#2117)

2020-01-29 12:29:03 -05:00

flatten_util.py

Remove "from __future__" uses from JAX. (#2117)

2020-01-29 12:29:03 -05:00

lax_linalg.py

remove input shapes from params of some primitives (#2410)

2020-03-13 07:13:29 -07:00

lax_reference.py

Fix test failures in incomplete gamma functions with Jaxlib 0.1.38. (#2118)

2020-01-29 13:24:45 -05:00

lazy.py

Remove "from __future__" uses from JAX. (#2117)

2020-01-29 12:29:03 -05:00

linear_util.py

Added some pytype annotations (#2386)

2020-03-09 20:41:01 +01:00

pprint_util.py

Remove "from __future__" uses from JAX. (#2117)

2020-01-29 12:29:03 -05:00

profiler.py

Add a jax.profiler module that exposes the new TensorFlow profiler in… (#2354)

2020-03-05 12:07:57 -05:00

random.py

Fix exception causes all over the codebase (#2376)

2020-03-09 16:06:12 -04:00

test_util.py

Allow zero tolerance for jax.test_util.tolerance (#2393)

2020-03-11 13:19:46 -07:00

tree_util.py

add register_pytree_node_class, fixes #2396 (#2400)

2020-03-10 15:01:18 -07:00

util.py

Removed unused imports (#2385)

2020-03-09 20:42:08 +01:00

version.py

bump version and update changelog for pypi

2020-02-11 07:22:17 -08:00