364 Commits

Author SHA1 Message Date
Jake VanderPlas
3520e4ddd4 doc: note in jax.scipy.special.erf that complex inputs are not supported. 2022-10-12 10:02:47 -07:00
Jake VanderPlas
afe74b4710 [typing] add type annotations to jax.scipy.linalg 2022-10-10 16:54:29 -07:00
Jake VanderPlas
124021d720 [typing] add annotations to jax.scipy.ndimage 2022-10-10 09:11:13 -07:00
Jake VanderPlas
76e4a1d40f [typing] add annotations to jax.scipy.fft 2022-10-09 05:18:45 -07:00
Jake VanderPlas
730f9930aa Update scipy version in jax.scipy.fft 2022-10-07 09:57:39 -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
Wonhyeong Seo
3f6eb40698 JAX implementation of scipy.stats.multinomial pmf & logpmf
Co-authored-by: harryjulian <harry.julian@peak.ai>
2022-09-15 13:21:44 -07:00
Peter Hawkins
57b5acf1b6 Roll forward: Upgrade logistic into a primitive.
Unlike the previous attempt, we don't try to use mhlo.logistic as the lowering of the new primitive yet. Instead, we lower to the old implementation of `expit`. This means that this change should be a no-op numerically and we can work on changing its implementation in a subsequent change.

PiperOrigin-RevId: 472705623
2022-09-07 06:06:56 -07:00
jax authors
9c16c83234 Rollback of upgrade logistic (sigmoid) function into a lax primitive.
PiperOrigin-RevId: 471105650
2022-08-30 15:30:43 -07:00
Peter Hawkins
f68f1c0cd0 Upgrade logistic (sigmoid) function into a lax primitive.
This allows us to lower it to `mhlo.logistic`, which allows XLA to generate more efficient code.

PiperOrigin-RevId: 470300985
2022-08-26 11:58:28 -07:00
jax authors
3e3542b0d6 Upgrade logistic (sigmoid) function into a lax primitive.
This allows us to lower it to `mhlo.logistic`, which allows XLA to generate more efficient code.

PiperOrigin-RevId: 469841487
2022-08-24 15:39:37 -07:00
Peter Hawkins
6276194e1c Upgrade logistic (sigmoid) function into a lax primitive.
This allows us to lower it to `mhlo.logistic`, which allows XLA to generate more efficient code.

PiperOrigin-RevId: 469789339
2022-08-24 12:04:01 -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
jax authors
0b6657e471 Merge pull request #11556 from RuffaloLavoisier:tYpO
PiperOrigin-RevId: 462648717
2022-07-22 10:13:10 -07:00
Jake VanderPlas
9769a0accf DOC: ensure that _wraps() generates correct links to wrapped functions 2022-07-21 11:12:35 -07:00
Jake VanderPlas
114b03670c Add missing f-string marker 2022-07-20 10:48:07 -07:00
RuffaloLavoisier
9f770425ac Correct spelling on word 2022-07-20 18:57:12 +09:00
Jake VanderPlas
9090dd179d jax.scipy.linalg.solve: deprecate the sym_pos argument following scipy 1.9.0 2022-07-19 13:57:49 -07:00
Bart van Merriënboer
de08344cb7 Avoid casting input to _fft_helper. 2022-07-06 14:29:54 -04:00
Jake VanderPlas
93deb86710 jax.scipy.optimize: fix type inconsistency 2022-06-29 09:47:53 -07:00
Dan F-M
0788d5708a Implementation of jax.scipy.stats.gaussian_kde 2022-06-28 15:17:12 -04:00
Peter Hawkins
a560a29e12 Increase the minimum scipy version to 1.5.
We don't have a formal support policy for scipy versions, but 1.5 dates from around the same date as the oldest supported NumPy release NEP-29 would have us support (1.20).
2022-06-24 15:07:09 -04:00
jax authors
15a3798424 Merge pull request #11147 from jakevdp:x64-scipy-signal-test
PiperOrigin-RevId: 456389803
2022-06-21 17:32:15 -07:00
Jake VanderPlas
3a8f478b0a [x64] make scipy_optimize_test compatible with strict dtype promotion 2022-06-21 09:29:08 -07:00
Jake VanderPlas
03f2189f90 [x64] make jax.scipy.signal compatible with strict dtype promotion.
Also a fair bit of cleanup & refactoring of related code.
2022-06-21 09:28:46 -07:00
Jake VanderPlas
4c4a83b108 [x64] make jax.scipy.sparse.linalg compatible with strict dtype promotion 2022-06-17 14:04:05 -07:00
Jake VanderPlas
1567adfa47 BUG: fix incorrect order of operations in line search 2022-06-16 15:21:24 -07:00
Jake VanderPlas
b5ba210097 [x64] make linalg functions & tests compatible with strict dtype promotion 2022-06-16 10:32:20 -07:00
jax authors
062a557476 Merge pull request #11100 from jakevdp:x64-scipy-stats-test
PiperOrigin-RevId: 455259751
2022-06-15 17:50:07 -07:00
Jake VanderPlas
297a2969a5 [x64] make fft functionality compatible with strict dtype promotion 2022-06-15 10:10:44 -07:00
Jake VanderPlas
f00d706a6d [x64] make scipy_stats_test.py compatible with strict dtype promotion 2022-06-14 14:47:58 -07:00
jax authors
d418e891f3 Merge pull request #11052 from jakevdp:x64-lax-scipy-test
PiperOrigin-RevId: 454940487
2022-06-14 13:22:53 -07:00
carlosgmartin
57b89ba7cb Added scipy.stats.gennorm. 2022-06-14 13:38:24 -04:00
Jake VanderPlas
e888e7c10c [x64] make lax_scipy_test.py compatible with strict dtype promotion 2022-06-14 10:02:45 -07:00
Jake VanderPlas
ebd53a48ab [x64] gmres: avoid problematic type promotions 2022-06-09 14:05:56 -07:00
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