mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-29 20:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
14,385 Commits 508 Branches 263 Tags
Commit Graph

113 Commits

Author SHA1 Message Date
Jake VanderPlas
924894fdd6 [x64] make tests more type-safe 2022-12-02 13:21:35 -08:00
Ian Horn
a35fe206a1 Added more accurate version of the betaln function. 2022-11-29 11:56:07 -08:00
Peter Hawkins
e42e52d4aa Rename test flag --num_generated_cases to --jax_num_generated_cases. 
parse_flags_with_absl() only parses flags that start with --jax_. Other flags are only parsed when absl.app's main function runs. But that's too late for test cases: test cases need to have the number of generated cases chosen at module initialization time. Hence the --num_generated_cases flag wasn't doing anything. Oops. By renaming it it works once again.

It might make sense to stop using flags for the number of generated cases and only use environment variables. We defer that to a future change.

Fix many test cases that were shown to be broken with a larger number of test cases enabled.

PiperOrigin-RevId: 487406670
 2022-11-09 18:58:05 -08:00
Tianjian Lu
3b1ddf2881 [linalg] Add jax.scipy.special.bessel_jn (Bessel function of the first kind). 
PiperOrigin-RevId: 487146250
 2022-11-08 23:03:21 -08:00
Peter Hawkins
72f4f389be Migrate remaining tests from jtu.cases_from_list to jtu.sample_product. 
Delete jtu.cases_from_list.
 2022-10-12 15:20:53 +00:00
Peter Hawkins
9bb2c999d6 Reenable some tests disabled in the past because of an LLVM bug. 
The issue no longer reproduces at head.

PiperOrigin-RevId: 480505525
 2022-10-11 18:59:37 -07:00
jax authors
363cc124e3 Merge pull request #12197 from ROCmSoftwarePlatform:fixedRocmUnitTestsSkip 
PiperOrigin-RevId: 479566021
 2022-10-07 06:36:11 -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
Chao Chen
1c69f594fb testSphHarmOrderZeroDegreeOne and test_custom_linear_solve_cholesky have been fixed in ROCm, no need to skip 2022-09-01 13:27:23 +00:00
Peter Hawkins
335b2cfb26 [JAX] Prepare not to export jax._src by default. 
Currently
```
import jax
```
populates `jax._src` in the names exported from JAX. This change prepares for not exporting `jax._src` by default.

In particular, explicitly import modules from jax._src and refer to those imports rather than assuming jax._src contents will be around later. This is a common pattern in tests.

This change does not yet remove any exported names.

Issue https://github.com/google/jax/issues/11951

PiperOrigin-RevId: 469480816
 2022-08-23 09:36:47 -07:00
Matthew Johnson
68e3f58041 un-skip polar/qdwh decomp tests skipped on gpu in ad6ce74 
On an A100 machine, these tests seem to run fine now. See https://github.com/google/jax/issues/8628#issuecomment-1215651697.
 2022-08-15 12:31:43 -07:00
Jake VanderPlas
7cc6b4f62b Tests: remove obsolete dtype_promotion decorators 2022-08-11 14:31:30 -07:00
Jake VanderPlas
e888e7c10c [x64] make lax_scipy_test.py compatible with strict dtype promotion 2022-06-14 10:02:45 -07:00
Peter Hawkins
400612f751 Relax test tolerance for spectral divide and conquer eigh test. 
Fixes test failures due to an overly tight tolerance.
 2022-05-31 09:36:07 -04:00
Peter Hawkins
045993411d Extend QDWH-eig implementation to work for complex inputs. 
We need to take only the real part when estimating the median Eigenvalue for splitting, and to return only the real part at termination.

PiperOrigin-RevId: 451406106
 2022-05-27 08:55:04 -07:00
Jake VanderPlas
5f7cd72130 [x64] use explicit casting rules for promote_dtypes_inexact 2022-05-24 15:51:44 -07:00
Peter Hawkins
b64e36b60f Make QDWH-eig implementation jit-table. 
Move QDWH-eig from jax._src.scipy.eigh to jax._src.lax.eigh, in preparation for using it to back `lax.eigh` in a future change.

PiperOrigin-RevId: 449362382
 2022-05-17 18:12:48 -07:00
Peter Hawkins
7ba36fc178 Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh. 
Remove jax._src.lax.polar.

PiperOrigin-RevId: 448241206
 2022-05-12 07:20:52 -07:00
Tianjian Lu
48f47c36c4 [linalg] Fix a bug in computing derivatives of scipy.special.lpmn. 
PiperOrigin-RevId: 447807140
 2022-05-10 13:02:58 -07:00
YouJiacheng
b485b8e5ce implement scipy.cluster.vq.vq 
also add no check_finite and overwrite_* docstring for some scipy.linalg functions
 2022-04-23 03:14:32 +08:00
Reza Rahimi
a0d9d81f92 Update JAX to use new math libraries in ROCm-5.0. 2022-03-01 20:02:15 +00:00
Jake VanderPlas
97512e9e44 JaxTestCase: set jax_numpy_rank_promotion='raise' by default 2022-02-14 09:22:05 -08:00
jax authors
5691010d2f Copybara import of the project: 
--
d42fffd849a4bac0c0c11a3346c93f07f8c64c44 by Jake VanderPlas <jakevdp@google.com>:

JaxTestCase: set numpy_rank_promotion='raise' by default
PiperOrigin-RevId: 427896974
 2022-02-10 19:08:29 -08:00
Jake VanderPlas
6324577a63 JaxTestCase: set numpy_rank_promotion='raise' by default 2022-02-10 16:54:31 -08:00
Jake VanderPlas
df0969961b Testing: avoid hard-coding random seeds 2021-12-10 10:32:09 -08:00
Diego Caballero
4f80e6af6e [XLA] More error tolerance adjustments due to incoming change in XLA:CPU 
We are about to change the vectorization strategy for XLA:CPU. As a result, we have
to increase the error tolerance of the impacted tests.

PiperOrigin-RevId: 412872724
 2021-11-29 06:26:55 -08:00
Diego Caballero
4d64677277 [XLA] Adjust the error tolerance of tests impacted by upcoming change in XLA:CPU 
We are about to change the vectorization strategy for XLA:CPU. This change may lead
to some numerical differences due to the fact the vectorization might happen differently
(e.g., code that was scalar could now be vectorized, code that was vectorized could now
be scalar, vectorization may happen with a different VL, etc.). As a result, we have
to increase the error tolerance of the impacted tests.

PiperOrigin-RevId: 412061380
 2021-11-24 08:08:04 -08:00
Peter Hawkins
ad6ce74d67 Skip some polar decomposition tests that fail on A100. 
Works around https://github.com/google/jax/issues/8628

PiperOrigin-RevId: 411604717
 2021-11-22 11:18:22 -08:00
jax authors
3a37cb1eb9 Merge pull request #8291 from jakevdp:lpmn-tests 
PiperOrigin-RevId: 404601366
 2021-10-20 11:35:50 -07:00
Aden Grue
349d0d0879 Relax the tolerances in one 'lax_scipy_test' case 
An upcoming CPU codegen change increases error slightly above the current tolerances.

PiperOrigin-RevId: 404402366
 2021-10-19 17:13:30 -07:00
Jake VanderPlas
58cdc1b3d6 special.lpmn: use more canonical testing approach 2021-10-19 13:46:32 -07:00
Aden Grue
d697ce2047 Relax the tolerances in one 'lax_scipy_test' case 
An upcoming CPU codegen change increases error slightly above the current tolerances.

PiperOrigin-RevId: 399570540
 2021-09-28 17:53:51 -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
Peter Hawkins
9f083d11da Use jax.* APIs rather than api.* names in tests. 
Tests should use our own public APIs where they exist.
 2021-09-13 16:01:32 -04:00
Jonathan Terhorst
fec72e1852 add support for scipy.special.{expn,expi,exp1} 2021-08-24 16:36:10 -04:00
Jake VanderPlas
730ae33e03 logsumexp: fix issue with debug_nans 2021-08-18 13:57:00 -07:00
Jake VanderPlas
0e256ddeb7 Fix logsumexp issue with debug_nans and disable_jit 2021-08-17 13:47:06 -07:00
Jake VanderPlas
6114e6a0d3 test_util: add decorator to set config values in test cases 2021-08-05 14:06:37 -07:00
Jake VanderPlas
30ea76cb6b disable rank promotion for jax scipy tests 2021-08-04 10:44:23 -07:00
Jake VanderPlas
20cef7eaa8 Fix rank promotion error in jsp.special.zeta 2021-08-03 14:33:57 -07:00
Jake VanderPlas
af16177659 Fix rank promotion error in jsp.special.multigammaln 2021-08-03 13:39:29 -07:00
Jake VanderPlas
e131343274 Fix issue with infinities in logsumexp 2021-08-02 15:27:24 -07:00
Peter Hawkins
0dfd76af97 Remove additional info return value from jax.scipy.linalg.polar(). 2021-07-20 13:13:31 -04:00
Adam Lewis
a2073ffcc2 Adds an implementation of a QR-based Dynamically Weighted Halley iteration. 2021-07-20 11:30:36 -04:00
tlu7
d97b393694 Adds spherical harmonics. 
Co-authored-by: Jake VanderPlas <jakevdp@google.com>
 2021-07-02 10:42:29 -07:00
tlu7
095e6507b9 Support value computation of associated Legendre functions. 
Co-authored-by: Jake VanderPlas <jakevdp@google.com>
 2021-06-14 14:51:37 -07:00
tlu7
a02bf59233 Adds associated Legendre functions of the first kind. 
Co-authored-by: Jake VanderPlas <jakevdp@google.com>
 2021-06-02 11:37:37 -07:00
Filippo Vicentini
c0c8e0d0a3 make logsumexp work with complex numbers 2021-05-31 16:01:57 +02:00
Peter Hawkins
26e9ebcdae Move jax.api to jax._src.api. 
PiperOrigin-RevId: 368233837
 2021-04-13 09:43:24 -07:00
Jake VanderPlas
7580876e31 jax.scipy.special.logsumexp: support integer input 2021-03-04 14:41:30 -08:00