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
12,579 Commits 508 Branches 263 Tags
Commit Graph

595 Commits

Author SHA1 Message Date
Jake VanderPlas
1a327f1ddd jnp.linalg.matrix_rank: support stacks of matrices 2022-08-08 11:16:34 -07:00
jax authors
33454f3097 Merge pull request #11762 from cglwn:add-inf-norm-error-for-string-ord 
PiperOrigin-RevId: 465764309
 2022-08-06 08:05:08 -07:00
Anthony
736bed0344 Fix jax numpy norm bug where passing ord=inf always returns one 2022-08-05 11:37:31 -07:00
Penn
1987ca7389 Add dtype arg to jnp.concatenate and update tests 2022-08-01 15:48:40 -07:00
jax authors
4fd700293e Merge pull request #11581 from jakevdp:fix-iscomplexobj 
PiperOrigin-RevId: 463184295
 2022-07-25 14:56:01 -07:00
jax authors
0b6657e471 Merge pull request #11556 from RuffaloLavoisier:tYpO 
PiperOrigin-RevId: 462648717
 2022-07-22 10:13:10 -07:00
Jake VanderPlas
4a693400b9 BUG: make jnp.iscomplexobj compatible with jit 2022-07-21 16:56:29 -07:00
jax authors
540ee56ff2 Merge pull request #11576 from jakevdp:searchsorted-alt 
PiperOrigin-RevId: 462461853
 2022-07-21 13:47:43 -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
10411bfeae jnp.searchsorted: add optional method argument to control implementation 2022-07-21 09:40:18 -07:00
RuffaloLavoisier
9f770425ac Correct spelling on word 2022-07-20 18:57:12 +09:00
Jake VanderPlas
2f4c485a54 Add dlpack support to device_array and jax.numpy 2022-07-15 17:31:11 -07:00
jax authors
a24100d5cf Merge pull request #11447 from jakevdp:empty-impl 
PiperOrigin-RevId: 460543062
 2022-07-12 13:45:28 -07:00
Jake VanderPlas
833a6ed450 jnp.ndarray.clip: remove deprecated arguments 2022-07-11 16:02:31 -07:00
jax authors
183b9e4503 Merge pull request #11397 from jakevdp:diagonal-err 
PiperOrigin-RevId: 460236416
 2022-07-11 09:52:19 -07:00
Jake VanderPlas
6c478edcef Improve documentation of jnp.empty() and jnp.empty_like() 2022-07-11 09:21:39 -07:00
Jake VanderPlas
e19df1a9bf Use asarray rather than array in ScalarMeta 
Why? This will make it so that jnp.int32(x) and friends no longer insert
a gratuitous copy_p operation in the jaxpr
 2022-07-08 11:16:40 -07:00
Jake VanderPlas
17de5e4840 jnp.diagonal: raise explicit error if ndim < 2 2022-07-07 16:36:40 -07:00
Marc van Zee
9d18f43a01 Do not normalize FFT by a constant "1" if no normalization is provided (i.e., norm is None). 
Without this, the compiled graph will still contain a node multipying a complex number with a constant 1+0j (1 is cast to complex because the other term is complex as well). This is problematic when converting to TFLite using jax2tf, because multiplying complex numbers is not supported in TFLite. With this change, the multiplication is removed from the graph all together.

PiperOrigin-RevId: 459566727
 2022-07-07 11:54:39 -07:00
Matthew Johnson
98e71fe31d [dynamic-shapes] revive basic bounded int machinery, add tests 2022-07-06 22:31:26 -07:00
Matthew Johnson
6bb90fde9e [dynamic shapes] revive iree 2022-07-06 15:01:16 -07:00
George Necula
5983d385da [dynamic-shapes] Expand the handling of dynamic shapes for reshape and iota. 
Also add more tests.
 2022-07-05 12:14:15 +03:00
jax authors
33f1f40b20 Merge pull request #11298 from pschuh:axis-cache-env 
PiperOrigin-RevId: 458328457
 2022-06-30 15:42:48 -07:00
Parker Schuh
6c5d204d7e Jax caches should depend on axis env. 2022-06-29 14:25:14 -07:00
Jake VanderPlas
8336a8b2d8 DeviceArray: raise explicit NotImplementedError for arr.flat 2022-06-29 10:11:22 -07:00
Jake VanderPlas
39b0ff7eb6 jnp.ndarray: raise TypeError for binary operations with builtin collections 2022-06-29 08:22:05 -07:00
Jake VanderPlas
df800f39d3 jnp.average: support keepdims argument 2022-06-28 10:55:55 -07:00
Yash Katariya
e32373c3ea Make jnp.array return jax.Array. Add input and result handlers for jax.Array. Also added tests for add under jit. 
TODO:
* Don't allow `x + y` if `jax.Array` is not fully addressable.
* Figure out how to use the already written tests with Array. Might be able to follow the path taken by SDA.
PiperOrigin-RevId: 457034779
 2022-06-24 10:05:06 -07:00
Jake VanderPlas
f6476f7a03 jnp.roots: better support for computation under JIT 2022-06-23 14:48:53 -07:00
Jake VanderPlas
e92e23e5f8 Use equality rather than identity when checking for float0 
Why? This is required due to changes to dtype canonicalization in numpy v1.23; see #11221
 2022-06-23 11:46:20 -07:00
Jake VanderPlas
15a19969de [x64] make polynomial_test compatible with strict dtype promotion 2022-06-17 16:52:22 -07:00
Jake VanderPlas
b5ba210097 [x64] make linalg functions & tests compatible with strict dtype promotion 2022-06-16 10:32:20 -07:00
Jake VanderPlas
297a2969a5 [x64] make fft functionality compatible with strict dtype promotion 2022-06-15 10:10:44 -07:00
Pavel Sountsov
ff637e12f1 Allow doing reductions on empty arrays in some cases. 
Namely, when the reduction axis is not over the zero-sized dimension.
 2022-06-14 21:57:56 +00:00
Jake VanderPlas
e8690f6ba3 [x64] preserve weak types in promote_dtypes_inexact 2022-06-14 09:34:43 -07:00
Jake VanderPlas
06c3857321 jnp.unique: improve error when run under JIT 2022-06-08 15:57:41 -07:00
Jake VanderPlas
412cc88a56 [x64] make jnp.modf() compatible with strict dtype promotion 2022-06-03 15:23:06 -07:00
Jake VanderPlas
010e490128 [x64] make jax.numpy reductions respect input dtypes 
Also make then compatible with strict dtype promotion mode.
 2022-06-01 16:24:36 -07:00
jax authors
7bb367b259 Merge pull request #10936 from jakevdp:x64-average 
PiperOrigin-RevId: 452401746
 2022-06-01 15:43:52 -07:00
jax authors
d762cf4511 Merge pull request #10939 from jakevdp:x64-insert 
PiperOrigin-RevId: 452400609
 2022-06-01 15:37:59 -07:00
jax authors
fe2968c58c Merge pull request #10938 from jakevdp:x64-corrcoef 
PiperOrigin-RevId: 452394751
 2022-06-01 15:10:03 -07:00
Jake VanderPlas
ed4962162d [x64] make jnp.insert compatible with strict dtype promotion 2022-06-01 15:00:53 -07:00
jax authors
e9542bb61d Merge pull request #10935 from jakevdp:x64-linspace 
PiperOrigin-RevId: 452388829
 2022-06-01 14:44:35 -07:00
Jake VanderPlas
92b0677a4e [x64] make jnp.corrcoef compatible with strict dtype promotion 2022-06-01 14:33:11 -07:00
Jake VanderPlas
1c555dc956 [x64] make jnp.average compatible with strict promotion 2022-06-01 14:25:35 -07:00
Jake VanderPlas
b916f07fa8 [x64] make linspace, logspace, & geomspace compatible with strict promotion mode 2022-06-01 14:09:04 -07:00
Jake VanderPlas
f710ec31e4 [x64] make jnp.interp safe for use with strict dtype promotion 2022-06-01 13:38:06 -07:00
Jake VanderPlas
12900bf6ab [x64] make jnp.mgrid compatible with strict dtype promotion 2022-06-01 10:21:49 -07:00
Jake VanderPlas
358f929681 [x64] jnp.ldexp: avoid implicit 64-bit promotion 2022-06-01 09:14:47 -07:00
jax authors
a1f7ced537 Merge pull request #10904 from jakevdp:x64-trapz 
PiperOrigin-RevId: 452205618
 2022-05-31 21:01:20 -07:00