There are a few test cases that generate millions of configurations,
only to have a handful of them selected by `cases_form_list`. I've
found all tests that spend over 100ms in case generation and
converted them to a new "test sampler" approach. The result: test
generation time drops from 15s to around 2s. Doesn't sound like much,
but I expect that we all run tests many times daily, so it seems like a
useful thing to have.
The rough idea is that the sampling generators get parameterized by a
sampler function that should be applied to the range of every `for` loop.
This allows us to sample runs of the generator through different
configurations by restricting each loop to a smaller subset. Right now
we always narrow it down to a single randomly selected instance. But,
we still retain the possibility of adding exhaustive testing in the
future, which can be achieved by passing in an identity sampling
function that wouldn't modify any loop ranges.
Rather than merely reporting a failure in check_grads(), we now
report the *specific* check that failed, e.g., "JVP tangent" or
"VJP of JVP cotangent projection". Gradient tests often fail for
spurious reasons (e.g., due to insufficient precision), so this should
be helpful for debugging.
I tested this manually by relaxing the tolerance for a test in
`linalg_test.py`.
Also, pass the body to XLA JIT when no parallel resources are used.
There is no reason to not do that given that we already require users to
pay the price of making their code jittable.
* Adds support for jit of pmap and pmap of pmap.
* Also adds a `tap_with_device` optional argument to `id_print` and
`id_tap`, to have the tap function invoked with a device keyword argument.
* Added multiple tests involving pmap
Issue: #5134Fixes: #5169
AD didn't use `HashableFunction` enough, tripping up the compilation
cache. I've also used the occasion to make function hashing a little
safer by including the Python bytecode of the wrapped function as part
of the key.
Changes:
- Fix unnecessary generator
- Iterate dictionary directly instead of calling .keys()
- Remove global statement at the module level
- Use list() instead of a list comprehension
- Use with statement to open the file
- Merge isinstance calls
* Add jax.linear_transpose
Co-authored-by: Matthew Johnson <mattjj@google.com>
* add failing test for complex numbers
* Add picky dtype check for linear_transpose
* Lint fix
* Allow truncating dtypes to match inputs in linear_transpose
* Fix typo in shape check error
* improve docstring
* Don't support integer inputs; better docstring
* fixup
* Fix doctest
Co-authored-by: Matthew Johnson <mattjj@google.com>
* Simplify the internal interface for host_callback.id_tap
This is a breaking change for `id_tap` users (but not `id_print` users).
This makes it easier to use (and type check) ``tap_func``, because the
expected signature is now ``tap_func(arg, transforms)`` vs
``tap_func(arg, *, transforms, **kwargs)``.
Most of the test changes are just adding whitespace/indentation, but I've
also slightly changed the way transformations are printed.
* 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.
