25 Commits

Author SHA1 Message Date
Jake VanderPlas
f090074d86 Avoid 'from jax import config' imports
In some environments this appears to import the config module rather than
the config object.
2024-04-11 13:23:27 -07:00
Peter Hawkins
6be860bda8 Clean up some device opt-in/opt-outs in test suite.
Use allowlists rather than denylists in a few places.

PiperOrigin-RevId: 568968749
2023-09-27 14:56:00 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
Peter Hawkins
c657449528 Copybara import of the project:
--
d39bdefb33a19e407c352df27fb04127f4fe8a1d by Peter Hawkins <phawkins@google.com>:

Migrate more tests from jtu.cases_from_list to jtu.sample_product.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/12717 from hawkinsp:sampletest d39bdefb33a19e407c352df27fb04127f4fe8a1d
PiperOrigin-RevId: 480136538
2022-10-10 11:35:32 -07:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
2022-09-22 12:27:19 -07:00
Peter Hawkins
29d03160e3 Remove _ prefix from functions in jax._src.dtypes.
to_inexact_dtype and to_complex_dtype are used across the JAX code base,
so they shouldn't have _ prefixes.
2022-08-12 12:51:09 +00:00
Jake VanderPlas
f6476f7a03 jnp.roots: better support for computation under JIT 2022-06-23 14:48:53 -07:00
Jake VanderPlas
df0969961b Testing: avoid hard-coding random seeds 2021-12-10 10:32:09 -08:00
Peter Hawkins
256e7220ff [JAX] Fix pylint errors.
* trailing-whitespace
* dangerous-default-value. None of these appear to be bugs in practice, but the potential for accidentally mutating the default value is there, and the cost of avoiding the problem is small.
* invalid-envvar-default. Pass strings as getenv() defaults.
* unnecessary-semicolon. Use tuples instead for this one-liner.
* invalid-hash-returned. Raise an exception rather than asserting false.
* pointless-string-statement. Use comments instead.
* unreachable. Use @unittest.skip() decorator rather than raising as first line in test.
* logging-not-lazy. Make the logging lazy.
* bad-format-string-type. Use f-string instead.
* subprocess-run-check. Pass check=...

PiperOrigin-RevId: 400858477
2021-10-04 17:54:46 -07:00
Peter Hawkins
db2e91eba2 Move jax.test_util to jax._src.test_util.
Add forwarding shims for names used by external clients of JAX in practice.

PiperOrigin-RevId: 398721725
2021-09-24 07:02:49 -07:00
Skye Wanderman-Milne
665312e355 Skip unimplemented TPU cases in polynomial_test.py 2021-03-04 00:12:55 +00:00
Jake VanderPlas
1a83bb6f90 Cleanup: remove remaining instances of rng_factory boilerplate 2020-12-11 13:47:46 -08:00
Peter Hawkins
c5b260c7f8 Relax some test tolerances and disable some tests on GPU. 2020-12-01 16:02:06 -05:00
Peter Hawkins
aa107cf1f4 Move jax.numpy internals into jax._src.numpy. 2020-10-16 20:35:19 -04:00
Matthew Johnson
552d4bdc12 add todo 2020-10-16 13:48:38 -07:00
Matthew Johnson
e51163af32 only pass hashable values as static args 2020-10-16 13:11:56 -07:00
Julius Kunze
4d99caa222
Fix polynomial tests (#3933) 2020-08-01 08:37:31 -07:00
Jake Vanderplas
0a6b715cd4
Add _NOT_IMPLEMENTED attribute to jax.numpy (fixes #3689) (#3698) 2020-07-09 16:31:08 -07:00
Jake Vanderplas
19adce595c
Cleanup: use test_util dtypes where possible (#3695)
* Cleanup: use test_util dtypes where possible

* fix issue in fft test

* fix duplicate test name issue
2020-07-08 13:21:48 -07:00
Jake VanderPlas
afce718eb1 Add ability to specify individual test targets 2020-06-29 11:08:57 -07:00
Peter Hawkins
66cea0277c
Fix test failures on GPU. (#3572) 2020-06-26 12:50:22 -04:00
Peter Hawkins
0ea22b7e19
Use a whitelist to restrict visibility in top-level jax namespace. (#2982)
* Use a whitelist to restrict visibility in top-level jax namespace.

The goal of this change is to capture the way the world is (i.e., not break users), and separately we will work on fixing users to avoid accidentally-exported APIs.
2020-05-07 17:24:19 -04:00
Peter Hawkins
b1bc841ae5
Replace np -> jnp, onp -> np in more places. (#2973)
* Replace np -> jnp, onp -> np in more places.

Context: #2370

* Fix typo in random_test.py
2020-05-05 16:40:41 -04:00
Peter Hawkins
7116cc5b41
Improve JAX test PRNG APIs to fix correlations between test cases. (#2957)
* 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.
2020-05-04 23:00:20 -04:00
Mathis Gerdes
bd70db79ef
Port np.roots (#2250)
* Implement np.roots.

* Expose jit-compatible variant of np.roots.

General np.roots implementation has a value dependent output shape.
If the input coefficients are guaranteed to have no leading zeros,
output shape is independent of values. Skip checking for leading
zeros by setting a keyword argument.

* Fix typo.

* Make roots jit-argument keyword only.

Co-Authored-By: Stephan Hoyer <shoyer@google.com>

* Format docstring to enable parsing.

Co-Authored-By: Stephan Hoyer <shoyer@google.com>

* Add np.roots function to documentation.

* Add more tests for np.roots function.

- Include length 0 polynomial coefficients
- Test strip_zeros=False argument
- Test jit compiled version (only on cpu due to eigvals)
- Confirm that adding leading zeros while skipping check
  for them results in nan's (expected behavior)

* Fix bug in np.roots test.

The polynomial with coefficents [0] never fails because the number of
roots is 0.

* Avoid bug in eigvals and adjust test accuracy.

The parameters of the test that was changed are non-essential
since they test for how the code behaves given invalid inputs.

The accuracy in comparing to the numpy result is changed because
the algorithm in those cases is slightly changed with respect to
the original numpy algorithm (to allow jit).

Co-authored-by: Stephan Hoyer <shoyer@google.com>
2020-04-09 23:16:53 -07:00