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

114 Commits

Author SHA1 Message Date
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
Peter Hawkins
41a0f6a682 Remove a stale comment. 2021-09-21 17:23:05 -04:00
jax authors
4bc6b27021 Merge pull request #7966 from jakevdp:faster-conv 
PiperOrigin-RevId: 398072150
 2021-09-21 13:33:33 -07:00
Peter Hawkins
1163e218e8 Attempt to land https://github.com/google/jax/pull/6400 again. 
This PR changes `jax.numpy.array()` to avoid creating any on-device arrays during tracing. As a consequence, calls to `jnp.array()` in a traced context, such as `jax.jit` will always be staged into the trace.

This change may break code that depends on the current (undocumented and unintentional) behavior of `jnp.array()` to perform shape or index calculations that must be known statically (at trace time). The workaround for such cases is to use classic NumPy to perform shape/index calculations.

PiperOrigin-RevId: 398008511
 2021-09-21 09:06:40 -07:00
Jake VanderPlas
d7e94b9eef convolutions: use flip() to clean up reverse-indexing 2021-09-21 08:49:32 -07:00
jax authors
e67d49b8a8 Merge pull request #7838 from khdlr:map_coordinates_reflect_mirror 
PiperOrigin-RevId: 397062070
 2021-09-16 05:58:44 -07:00
jax authors
5c58afa647 Merge pull request #7855 from cyprienc:scipy-stats-nbinom 
PiperOrigin-RevId: 396617949
 2021-09-14 09:49:50 -07:00
Cyprien
8c8f0a8c71 Feat: scipy.stats.nbinom implementation 
fix: increasing tolerance check for testNBinomLogPmf in scipy_stats_test.py
 2021-09-14 08:05:42 +01:00
Peter Hawkins
8b2123968a Switch internal users of jax.util.partial to use functools.partial. 2021-09-13 21:09:58 -04:00
Peter Hawkins
a84426cb8f Switch internal users of jax.ops.index_... to use x.at[x].set() APIs. 2021-09-13 19:48:29 -04:00
Peter Hawkins
80599c0821 Replace uses of jax.partial with functools.partial, in preparation for removing jax.partial. 
jax.partial is an alias for functools.partial, and functools.partial is a Python standard library API. There's no need for jax to export this function.

PiperOrigin-RevId: 396370975
 2021-09-13 09:16:19 -07:00
Konrad Heidler
5ed619afbb
Implement 'reflect' and 'mirror' padding modes for scipy.ndimage.map_coordinates 2021-09-07 16:05:35 +02: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
jax authors
3827a8b26a Merge pull request #7617 from juliuskunze:dct 
PiperOrigin-RevId: 391540164
 2021-08-18 09:06:02 -07:00
Jake VanderPlas
0e256ddeb7 Fix logsumexp issue with debug_nans and disable_jit 2021-08-17 13:47:06 -07:00