mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 04:16:08 +00:00
Code Issues Packages Projects Releases Wiki Activity
11,796 Commits 508 Branches 263 Tags
Commit Graph

129 Commits

Author SHA1 Message Date
Peter Hawkins
9be53caa43 Use util.unzipN() in more places instead of zip(*args). 2022-06-01 16:28:31 -04:00
Sergei Lebedev
c5d3ece6f5 MAINT Fixed new mypy errors 
mypy seems to handle lambdas and named functions differently. So, I had to
promote a few helpers to named functions to get them to type check.
 2022-05-23 20:21:00 +01:00
Jeppe Klitgaard
838a05329d feat: validate jit args 2022-05-18 21:54:47 +01:00
jax authors
6110be40dc Merge pull request #10678 from JeppeKlitgaard:precommit-pyupgrade 
PiperOrigin-RevId: 449561541
 2022-05-18 13:25:52 -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
Jeppe Klitgaard
17de89b16a feat: refactor code using pyupgrade 
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```

a
 2022-05-17 22:14:05 +01: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
Peter Hawkins
705e241409 Change non-array arguments to jax.lax.linalg functions to be keyword-only arguments. 
PiperOrigin-RevId: 448066207
 2022-05-11 13:06:54 -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
Jake VanderPlas
30fd817130 jax.scipy.sparse.linalg: support sparse matrices as operators 2022-05-05 10:33:08 -07:00
jax authors
39ac9faf61 Merge pull request #10419 from YouJiacheng:jsp.linalg-unused-params 
PiperOrigin-RevId: 444840248
 2022-04-27 06:27:09 -07:00
Jake VanderPlas
822b6aad3b jax.scipy.qr: fix return type for mode='r' 2022-04-26 11:26:56 -07:00
YouJiacheng
ecab8e00e8 try to improve docs for scipy.linalg with unused parameters 2022-04-26 14:56:26 +08:00
Jake VanderPlas
78adfbd969 Add missing __init__ file in jax._src.scipy.cluster 2022-04-25 15:29:38 -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
YouJiacheng
af7b94b110 Fix typo of #10381 
and add a basic regression test
 2022-04-21 02:17:09 +08:00
YouJiacheng
bb2682db6d remove numpy.linalg._promote_arg_dtypes 
in favor of numpy.util._promote_dtypes_inexact
 2022-04-21 00:23:56 +08:00
jax authors
e8ae9d4dbb Merge pull request #10220 from YouJiacheng:Fix#10219 
PiperOrigin-RevId: 441515789
 2022-04-13 10:34:32 -07:00
YouJiacheng
dad324d934 Fix#10247 2022-04-13 04:05:55 +08:00
YouJiacheng
4695dd919c Fix#10219 2022-04-13 04:04:11 +08:00
Alex Riley
869596fc2c Add jax.scipy.linalg.rsf2csf 2022-04-06 21:06:23 +01:00
Jake VanderPlas
df1ceaeeb1 Deprecate jax.tree_util.tree_multimap 2022-04-01 14:51:54 -07:00
Yotaro Kubo
a7fd751acf Add istft to jax.scipy.signal. 2022-04-01 14:28:53 +09:00
Yotaro Kubo
2e70177385 Fix a bug in fft helper appears when nperseg=1. 2022-03-28 14:51:54 +09:00
Jake VanderPlas
998d60dd07 DOC: clarify parameter types in cg/bicgstab 2022-03-23 08:35:25 -07:00
Jake VanderPlas
0a72adbd5e lax_numpy: factor out indexing tricks 2022-03-17 11:05:45 -07:00
Roy Frostig
8f93629e87 remove _convert_element_type from public jax.lax module 2022-03-09 18:46:38 -08:00
Roy Frostig
f7731bf959 remove _const from public jax.lax module 
Modify all internal call sites to use `jax._src.lax.lax._const`.
 2022-03-07 12:26:25 -08:00
Jake VanderPlas
00e040e514 cleanup: remove _constant_like in favor of lax._const 2022-03-02 09:13:58 -08:00
jax authors
54a6e4dad3 Merge pull request #9422 from yotarok:signal_stft 
PiperOrigin-RevId: 429377655
 2022-02-17 12:46:12 -08:00
Yotaro Kubo
e085370ec4 Add some functions for spectral analysis. 
This commit adds "stft", "csd", and "welch" functions in scipy.signal.
 2022-02-17 15:59:24 +09:00
Leello Tadesse Dadi
cb732323f3 adds jax.scipy.linalg.sqrtm 2022-02-16 22:33:47 +01:00
Leello Tadesse Dadi
514d8883ce adds jax.scipy.schur 2022-02-16 22:33:37 +01:00
Jake VanderPlas
27f285782b linalg_test: disable implicit rank promotion 2022-01-26 09:29:06 -08:00
Attila Szabó
70bf281250 Fix max_squarings in expm 2022-01-25 09:54:23 +00:00
jax authors
16c809ce7f Merge pull request #8625 from Edenhofer:regular_grid_interpolator 
PiperOrigin-RevId: 423876494
 2022-01-24 12:02:34 -08:00
DanPuzzuoli
2d2ac12aa0 switching to use jnp.digitize for index identification 2022-01-19 10:02:06 -08:00
DanPuzzuoli
cbb2f7baab changed to use argwhere 2022-01-18 16:28:36 -08:00
DanPuzzuoli
9d43ccafa3 converting to use switch 2022-01-18 14:17:17 -08:00
Gordian Edenhofer
2c5fe8c40d Implement SciPy's RegularGridInterpolator 
Resolves #8572 .
 2022-01-04 22:10:36 +01:00
Jake VanderPlas
f8e18e9a00 [x64] minor weak_type changes to linalg.py 2021-12-07 16:27:29 -08:00
Jake VanderPlas
022f8ac2ee [x64] preserve weak types in jax.scipy.sparse solvers 2021-11-30 10:36:28 -08:00
jax authors
d81114deff Merge pull request #8668 from jakevdp:canonicalize-axis 
PiperOrigin-RevId: 412088801
 2021-11-24 10:30:14 -08:00
Peter Hawkins
4e21922055 Use imports relative to the jax package consistently, rather than .-relative imports. 
This is more consistent, since currently we use a mix of both styles. It may also help pytype yield more accurate types.

PiperOrigin-RevId: 412057514
 2021-11-24 07:48:29 -08:00
Jake VanderPlas
f6e3f1b4ad Cleanup: remove duplicate canonicalize_axis utility 2021-11-23 16:54:02 -08:00
Jake VanderPlas
f2a959054a Document jax.lax.Precision 2021-11-08 14:15:31 -08:00
Jake VanderPlas
40d6f5ed90 Tighten up dtypes across the package 2021-10-29 13:50:30 -07:00
Jake VanderPlas
0a232b2237 stats.multivariate_normal: support broadcasted inputs 2021-10-19 16:58:36 -07:00
Peter Hawkins
f8ba024621 Fix JAX functions to work if the default gather mode is set to "fill". 
These functions really do want "clip".
 2021-09-30 14:21:05 -04:00
Peter Hawkins
2eb20357db Add @jit decorators to jax.numpy.linalg and jax.scipy.linalg. 2021-09-24 15:52:11 -04:00