Adam Paszke adb442eb8a Make ad_util.zero a class that carries avals (similar to UndefinedPrimal) ...

This is useful for remat transpose rule submitted in #3162 and e.g.
allowed me to catch a slight overuse of defjvp2 for `random_gamma_p` (it
was unnecessarily declared as having multiple outputs).

2020-06-05 15:51:30 +00:00

..

experimental

Refactoring of jax_to_tf tests: (#3262) (#3308)

2020-06-04 09:54:06 +03:00

interpreters

Make ad_util.zero a class that carries avals (similar to UndefinedPrimal)

2020-06-05 15:51:30 +00:00

lax

Make ad_util.zero a class that carries avals (similar to UndefinedPrimal)

2020-06-05 15:51:30 +00:00

lib

Cleanup: deflake interpreters, lib, nn, third_party, and tools (#3327)

2020-06-04 15:27:48 -07:00

nn

Cleanup: deflake interpreters, lib, nn, third_party, and tools (#3327)

2020-06-04 15:27:48 -07:00

numpy

Improve error when zero-sized arrays passed to convolve (#3325)

2020-06-04 20:28:19 -07:00

ops

Update internal aliases to lax_numpy to jnp instead of np. (#2975)

2020-05-05 20:41:57 -04:00

scipy

Improve error when zero-sized arrays passed to convolve (#3325)

2020-06-04 20:28:19 -07:00

third_party

Cleanup: deflake interpreters, lib, nn, third_party, and tools (#3327)

2020-06-04 15:27:48 -07:00

tools

Update JAX to avoid XLA:Python API names deprecated in jaxlib 0.1.46. (#3046)

2020-05-11 17:43:55 -04:00

__init__.py

Exported lax from jax/__init__.py (#3135)

2020-05-19 15:40:03 -04:00

abstract_arrays.py

Change instances of onp to np and np to jnp (#3044)

2020-05-12 20:37:05 -04:00

ad_util.py

Make ad_util.zero a class that carries avals (similar to UndefinedPrimal)

2020-06-05 15:51:30 +00:00

api_util.py

Add support for buffer donation in jit and pmap. (#2936)

2020-05-31 15:00:16 -07:00

api.py

Make ad_util.zero a class that carries avals (similar to UndefinedPrimal)

2020-06-05 15:51:30 +00:00

BUILD

Implementation of id_tap/id_print using outfeed. (#3006)

2020-05-08 17:18:11 +03:00

config.py

Test code in docs and api.py docstrings (#2994)

2020-05-16 16:19:24 +03:00

core.py

change cond primitive to an indexed conditional with multiple branch functions

2020-06-03 22:19:15 -07:00

custom_derivatives.py

Make ad_util.zero a class that carries avals (similar to UndefinedPrimal)

2020-06-05 15:51:30 +00:00

dlpack.py

Update JAX to avoid XLA:Python API names deprecated in jaxlib 0.1.46. (#3046)

2020-05-11 17:43:55 -04:00

dtypes.py

Expand type support for random uniform() & randint() (#3138)

2020-05-19 14:19:00 -07:00

flatten_util.py

Replace np -> jnp, onp -> np in more places. (#2973)

2020-05-05 16:40:41 -04:00

lax_linalg.py

Make ad_util.zero a class that carries avals (similar to UndefinedPrimal)

2020-06-05 15:51:30 +00:00

lax_reference.py

Make check_dtypes, atol, and rtol keyword-only arguments in jax.test_… (#3280)

2020-06-01 17:19:23 -04:00

lazy.py

Change instances of onp to np and np to jnp (#3044)

2020-05-12 20:37:05 -04:00

linear_util.py

revamp custom_jvp/vjp implementation to fix bugs

2020-03-29 19:35:01 -07:00

pprint_util.py

change cond primitive to an indexed conditional with multiple branch functions

2020-06-03 22:19:15 -07: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

Make ad_util.zero a class that carries avals (similar to UndefinedPrimal)

2020-06-05 15:51:30 +00:00

test_util.py

Fix bug where jnp.array returned a classic NumPy array, sometimes wit… (#3283)

2020-06-01 19:29:26 -04:00

tree_util.py

Expose functools.reduce initializer argument to tree_util.tree_reduce (#2935)

2020-05-05 11:11:10 +03:00

util.py

merge constvars when forming cond branch jaxprs

2020-05-13 21:14:41 -07:00

version.py

update version and changelog for pypi

2020-06-03 07:45:53 -07:00