Matthew Johnson 42a1ad4307 change dtype promotion behavior for jit-invariance ...

Here are two desiderata for jax.numpy dtype promotion behavior:
1. follow what NumPy does
2. be invariant to `@jit`

The latter is much more important, so whenever the two are in tension we
prefer the latter. (Also we already can't do a perfect job following
what NumPy does, e.g. around its value-dependent dtype promotion logic.)

Issue #732 showed our code had a special behavior that essentially
handled a case of the former desideratum but also broke the latter. #732
also showed us (again) that our tests really should cover Python
scalars.

In summary, in this commit:
* revise jax.numpy dtype promotion behavior to be invariant to `@jit`
* add Python scalar types to lax_numpy tests
* simplify and update kron implementation to fix dtype issues

2019-05-19 18:49:16 -07:00

..

experimental

Adagrad optimizer

2019-05-13 20:36:45 -07:00

interpreters

catch pmap num_devices errors at compile time too

2019-05-17 13:56:39 -07:00

lax

change dtype promotion behavior for jit-invariance

2019-05-19 18:49:16 -07:00

lib

revise nested pmap gather to CRS-based impl

2019-05-17 10:25:26 -07:00

numpy

change dtype promotion behavior for jit-invariance

2019-05-19 18:49:16 -07:00

ops

Add segment_sum to the docs and fix its rendering.

2019-05-08 21:18:17 -04:00

scipy

change dtype promotion behavior for jit-invariance

2019-05-19 18:49:16 -07:00

__init__.py

Use a regular import to add jax.__version__ rather than exec() trickery.

2019-02-19 11:38:28 -05:00

abstract_arrays.py

Merge pull request #704 from google/differentiable-scan

2019-05-13 10:26:09 -07:00

ad_util.py

ship it

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

api_util.py

avoid packing leaf outputs for jit/pmap funs

2019-05-17 07:36:52 -07:00

api.py

avoid packing leaf outputs for jit/pmap funs

2019-05-17 07:36:52 -07:00

BUILD

Add lax/*.py to jax library dependencies.

2019-04-16 09:29:13 -04:00

config.py

Test case improvements:

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

core.py

add literals to jaxprs

2019-05-13 08:48:13 -07:00

flatten_util.py

prevent jit from treating keyword args as static

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

lax_linalg.py

Incorporate review comments.

2019-05-13 19:53:50 -04:00

lax_reference.py

source sync

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

linear_util.py

remove assert that fails on python3 map objects

2019-04-10 22:17:54 -07:00

pprint_util.py

source sync

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

random.py

fix randint docstring

2019-05-16 10:36:30 -07:00

test_util.py

change dtype promotion behavior for jit-invariance

2019-05-19 18:49:16 -07:00

tree_util.py

revies optimizers api, fix misc bugs

2019-05-03 12:44:52 -07:00

util.py

make static_argnums cache on value when possible

2019-05-09 20:00:24 -07:00

version.py

bump version number for pypi

2019-05-17 13:37:43 -07:00