* Fix eigh JVP to ensure that both the primal and tangents of the eigenvalues are real.
Add test to jax.test_util.check_jvp that ensure the primals and both the primals and tangents produced by a JVP rule have identical types.
* Cast input to static indexing grad tests to a JAX array so new type check passes.
* Fix bug where jnp.array returned a classic NumPy array, sometimes with the wrong type.
Unconditionally calls `device_put`, because `lax.convert_element_type` has a fast path that sometimes fails to lead to a `device_put`.
Improve the test for `jnp.array` and its test harness.
* Make check_dtypes, atol, and rtol keyword-only arguments in jax.test_util APIs.
Default to check_dtypes=True.
Remove explicit usages of check_dtypes=True from tests. This mostly just removes visual noise from tests. Testing for exact type equality is the sensible default, although there are cases where opting out makes sense.
No functional changes intended.
* Fix a number of lax reference implementations to preserve types.
* Improve JAX test PRNG APIs to fix correlations between test cases.
In #2863, we observed that we were missing gradient problems because the random test cases being generated were too similar because they were formed with identically seeded PRNGs. This change updates the test_util.rand_...() functions to take an explicit numpy.random.RandomState, and adds a rng() method to JaxTestCase to form a RandomState seeded on the test case name.
This gives the following properties:
* different test cases receive different seeds
* PRNG seeding is deterministic and independent of execution order and sharding.
* PRNG seeding is deterministic across runs.
* Fix some failing tests.
* Fix more test failures.
Simplify ediff1d implementation and make it more permissive when casting.
* Relax test tolerance of laplace CDF test.
A significant fraction of time when collecting test cases is spent building shape and dtype strings (which are usually similar and usually thrown away.)
* Make pytest run over JAX tests warning clean, and error on warnings.
Remove global warning suppression in travis.yml. Instead add a pytest.ini that converts warnings to errors, with the exception of a whitelist.
Either fix or locally suppress warnings in tests.
Also fix crashes on Mac related to a preexisting linear algebra bug.
* Fix some type errors in the FFT transpose rules revealed by the convert_element_type transpose rule change.
Introduced two new constructors for PartialVal: unknown and known.
These should make it easier to read the code where we construct
PartialVal:
* instead of PartialVal((aval, core.unit) we use PartialVal.unknown(aval)
* instead of PartialVal((None, pval)) we use PartialVal.known(pval)
Also disabled some new tests in random_tests.py on Mac. They segfault,
apparently due to the same issue #432.
fixes#2314
I also added a bit more test coverage, but not a ton: scipy has
different batch shape semantics and default arguments than I might
expect, so I didn't bother to implement those (and left some test cases
commented out).
I ran into this surprising scipy bug:
```python
In [1]: from scipy.stats import multivariate_normal
In [2]: import numpy as np
In [3]: args = [np.array(1., np.float32), np.array(2., np.float64), np.array(3., np.float64)]
In [4]: print([x.shape for x in args])
[(), (), ()]
In [5]: multivariate_normal.logpdf(*args)
Out[5]: -1.6349113442053944
In [6]: print([x.shape for x in args])
[(), (1,), (1, 1)]
```
Mutated arguments! But it depends on dtype promotion:
```python
In [7]: args = [np.array(1., np.float32), np.array(2., np.float32), np.array(3., np.float32)]
In [8]: print([x.shape for x in args])
[(), (), ()]
In [9]: multivariate_normal.logpdf(*args)
Out[9]: -1.6349113442053944
In [10]: print([x.shape for x in args])
[(), (), ()]
```
Currently, if a user passes any falsy value to jax.test_util.tolerance,
it is changed to the default value. This makes sense when the value
passed is None, but not when the value passed is 0 (which indicates
a desired tolerance of exactly 0).
Disables failing tests for now.
The goal is to make the Jaxpr language more uniform: all higher-order
primitives carry sub-Jaxprs that are part of the parameters, and they
are all called xxx_jaxpr. As a side-effect, some code is simplified
(e.g., the code that searches for sub-jaxprs).
For now the code assumes that all the `call` (final-style) primitives
carry exactly one subjaxpr with the parameter name `call_jaxpr`. These
primitives are still processed differently in the internal code, but
there is no reason any external consumer of a Jaxpr needs to know this.
Before, bound_subjaxprs was a tuple (0 or 1 values) of
a pair of a Jaxpr and its constant values. Now we close up all such Jaxprs
such that they do not take constvars and their constant values are part of the
arguments.
We also rename bound_subjaxprs to bound_subjaxpr (an optional Jaxpr)
This is first part of a simplification. In a subsequent PR I will move
the bound_subjaxpr into params, as for most higher-order primitives.
Freevars played a very small role, and they can be folded with
the invars. This simplifies the Jaxpr data structure.We remove
the `freevars` field from Jaxpr and from the bound_subjaxprs.
The only non-trivial change is for xla_pmap, where we need
to carry one extra parameter `mapped_invars` with a bitmap
to encode which invars are mapped and which are broadcast.
Previously, the freevars were broadcast.
This change prepares for switching the default types in JAX's NumPy to be 32-bit types. In particular, it makes the JAX tests pass in the event that jax.numpy.int_, jax.numpy.float_, and jax.numpy.complex_ are defined to be 32-bit types instead of 64-bit types, but does not yet change the defaults.
* Fixes to type handling.
* Specify exactly which types to test in lax_test.py, rather than relying on non-x64 mode to squash unsupported types.
* Fix some excessive promotions in jax.numpy.
* Fix some buggy RNGs that returned the wrong type for complex inputs.
* Remove np._promote_args_like, and replace its users with a newer _promote_args_inexact.
We no longer want to promote arguments exactly like NumPy; NumPy has a bad habit of promoting integer types to float64, whereas we want to promote to jax.numpy.float_, which may not be the same.
For example
```
import numpy as onp
onp.sin(3).dtype
```
returns `onp.dtype(float64)`.
However, it turns out that all of the users of `_promote_args_like` are using it for exactly one behavior: promoting integers or bools to inexact types like float. Implement that behavior explicitly rather than mimicing the behavior of NumPy.
* Relax test tolerances.
* A TypedJaxpr contains more useful information (consts, types)
* Also forced the instantiation of constants when producing the jaxpr.
Before:
>>>print(api.make_jaxpr(lambda x: 1.)(0.))
lambda ; ; a.
let
in [*]}
After this change:
>>>print(api.make_jaxpr(lambda x: 1.)(0.))
lambda ; ; a.
let
in [1.0]}
