Matthew Johnson d857d3f595 improve prng compile times with outer-loop rolling

2019-08-13 09:43:44 -07:00

..

experimental

Convert tableaux & runge_kutta_step to JAX control flow, use np.where

2019-08-12 22:11:42 +00:00

interpreters

Keep ShapedArray avals on xla.DeviceArray values

2019-08-12 10:03:04 -04:00

lax

Keep ShapedArray avals on xla.DeviceArray values

2019-08-12 10:03:04 -04:00

lib

More concurrency fixes.

2019-08-09 15:08:26 -04:00

numpy

Add test case for concurrent device_get and device_put calls.

2019-08-09 13:12:44 -04:00

ops

Remove debug print statement and stale comment.

2019-07-21 16:35:00 -04:00

scipy

Fix flake8 warnings, including a few real bugs.

2019-08-04 10:25:48 -04:00

tools

Replace jaxlib import paths with jax.lib paths.

2019-07-30 21:01:48 -04:00

__init__.py

import random in jax/__init__.py

2019-07-15 23:04:50 +01:00

abstract_arrays.py

Mark instances of 'long' with # noqa

2019-08-05 00:22:41 +02:00

ad_util.py

ship it

2019-05-08 13:41:25 -07:00

api_util.py

Python changes in preparation for adding a C++ implementation of the PyTree utilities.

2019-07-29 10:57:27 -04:00

api.py

make pmap read axis size from kwargs

2019-08-12 18:03:25 -07:00

BUILD

Remove old dummy build_jax target.

2019-06-27 22:13:04 -04:00

config.py

Test case improvements:

2019-04-12 10:48:11 -04:00

core.py

Make JAX tracer state thread-local. Allows performing traces in separate threads.

2019-08-09 13:55:20 -04:00

flatten_util.py

prevent jit from treating keyword args as static

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

lax_linalg.py

Check for symbolic zeros in linear algebra gradients.

2019-08-09 10:58:09 -04:00

lax_reference.py

source sync

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

linear_util.py

Add test case for concurrent device_get and device_put calls.

2019-08-09 13:12:44 -04:00

pprint_util.py

source sync

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

random.py

improve prng compile times with outer-loop rolling

2019-08-13 09:43:44 -07:00

test_util.py

Add a jtu.device_under_test() method. Use it instead of reparsing the device flag everywhere.

2019-08-04 17:17:49 -04:00

tree_util.py

Address review comments.

2019-08-01 16:48:18 -04:00

util.py

Add test case for concurrent device_get and device_put calls.

2019-08-09 13:12:44 -04:00

version.py

bump version for pypi

2019-08-08 09:34:03 -07:00