Skye Wanderman-Milne 8f4ba7e679

Allow specifying both devices and axis_size to pmap. (#3475) ...

This allows providing custom device assignments to nested pmaps or pmap-of-sharded_jit when running on a multi-host platform.

2020-06-19 15:51:12 -07:00

..

experimental

Fix test failure in jax2tf due to conflicting merges. (#3492)

2020-06-19 10:14:53 -04:00

interpreters

Allow specifying both devices and axis_size to pmap. (#3475)

2020-06-19 15:51:12 -07:00

lax

Add random_gamma_grad and use in jax.random.gamma (#3281)

2020-06-19 09:34:18 -04:00

lib

Attach source info to Jaxpr equations. (#3421)

2020-06-17 16:35:36 -07:00

nn

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

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

numpy

use original numpy for shape calculations (#3474)

2020-06-17 10:05:28 -07:00

ops

Cleanup: deflake jax.experimental and jax.ops (#3329)

2020-06-05 19:00:04 -07:00

scipy

add scipy.special.zeta, scipy.special.polygamma (#3385)

2020-06-15 19:21:26 -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

Initial version of invertible AD implementation (#3232)

2020-06-15 12:35:06 +02:00

abstract_arrays.py

Optimize zeros_like_shaped_array

2020-06-05 15:52:03 +00: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

deflake remainder of jax (#3343)

2020-06-06 10:51:34 -07:00

api.py

Allow specifying both devices and axis_size to pmap. (#3475)

2020-06-19 15:51:12 -07:00

BUILD

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

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

config.py

deflake remainder of jax (#3343)

2020-06-06 10:51:34 -07:00

core.py

Attach source info to Jaxpr equations. (#3421)

2020-06-17 16:35:36 -07:00

custom_derivatives.py

fix defjvps with None as first rule

2020-06-09 15:19:53 -07: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

Improve error message when passing an invalid dtype. (#3405)

2020-06-11 12:48:05 -07:00

flatten_util.py

deflake

2020-06-12 15:41:49 -07:00

lax_linalg.py

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

2020-06-08 17:50:14 +02:00

lax_reference.py

[jax2tf] Added special case for tf.pad. (#3462)

2020-06-17 11:57:21 +03:00

lazy.py

deflake remainder of jax (#3343)

2020-06-06 10:51:34 -07:00

linear_util.py

revamp custom_jvp/vjp implementation to fix bugs

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

pprint_util.py

Attach source info to Jaxpr equations. (#3421)

2020-06-17 16:35:36 -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

Add random_gamma_grad and use in jax.random.gamma (#3281)

2020-06-19 09:34:18 -04:00

source_info_util.py

Attach source info to Jaxpr equations. (#3421)

2020-06-17 16:35:36 -07:00

test_util.py

deflake remainder of jax (#3343)

2020-06-06 10:51:34 -07:00

tree_util.py

add docstring to ravel_pytree

2020-06-12 15:41:07 -07:00

util.py

merge constvars when forming cond branch jaxprs

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

version.py

Update jax version to 0.1.70 (#3383)

2020-06-09 09:55:05 -07:00