mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-21 22:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
8,296 Commits 508 Branches 263 Tags
Commit Graph

3028 Commits

Author SHA1 Message Date
jax authors
3b84f85700 Merge pull request #7058 from jakevdp:bcoo-serialization 
PiperOrigin-RevId: 380902758
 2021-06-22 15:13:31 -07:00
jax authors
986c390021 Merge pull request #7055 from hawkinsp:numpy 
PiperOrigin-RevId: 380892531
 2021-06-22 14:31:22 -07:00
jax authors
6f8ebb7f10 Merge pull request #6978 from colemanliyah:get_cache_key 
PiperOrigin-RevId: 380892453
 2021-06-22 14:27:27 -07:00
Jake VanderPlas
e18183eb45 [sparse] fix batched serialization of BCOO 2021-06-22 13:33:52 -07:00
Peter Hawkins
75c9bf01f3 Fix most test failures under NumPy 1.21. 2021-06-22 16:31:44 -04:00
Peter Hawkins
f885366a8b [JAX] Improve support for DLPack tensors on CPU when a GPU is available. 
https://github.com/google/jax/issues/5581

Previously the user had to provide the target backend explicitly. Now we supply both CPU and GPU backends to the C++ code so it can choose based on the metadata of the DLPack tensor.

PiperOrigin-RevId: 380795192
 2021-06-22 06:38:53 -07:00
Liyah Coleman
a87a291afd Implemented get_cache_key funtion with unit tests 2021-06-21 22:49:09 +00:00
Qiao Zhang
e8e6138c75 Add vmap rule for lax.clamp. 2021-06-21 14:59:34 -07:00
Jake VanderPlas
71bf410055 Fix random test for scipy>=1.17.0 2021-06-21 13:44:54 -07:00
jax authors
8bb6bdc4b8 Merge pull request #7025 from jakevdp:bcoo-batch 
PiperOrigin-RevId: 380561985
 2021-06-21 06:14:00 -07:00
Jake VanderPlas
1bf0829c21 [sparse] add tests of sparse dot_general autodiff with respect to sparse argument 2021-06-18 16:05:35 -07:00
Luke Pfister
c33388b136 Support complex numbers in jax.scipy.signal.convolve/correlate 2021-06-18 13:07:00 -06:00
jax authors
605389025c Merge pull request #6914 from njunge94:custom_root_aux 
PiperOrigin-RevId: 380227301
 2021-06-18 11:25:30 -07:00
Jake VanderPlas
a385d975b9 [sparse] add sparse-aware wrapper for grad & value_and_grad 2021-06-18 08:18:56 -07:00
Peter Hawkins
4406c77086 [XLA] Improve the accuracy of complex log1p close to 0. 
https://github.com/google/jax/issues/7004

PiperOrigin-RevId: 380188368
 2021-06-18 08:12:23 -07:00
George Necula
a4e74a269c Disable testing dot_general with preferred_element_type on GPU. 
Due to a XLA bug we get non-deterministic NaN on GPU.
This fixes flakiness in lax_test.py

PiperOrigin-RevId: 380179946
 2021-06-18 07:12:54 -07:00
George Necula
6a48c60a72 Rename master to main in embedded links. 
Tried to avoid the change on external links to repos that
have not yet renamed master.
 2021-06-18 10:00:01 +03:00
jax authors
1d0c461ea6 Merge pull request #6995 from jakevdp:sparse-submodule 
PiperOrigin-RevId: 380040999
 2021-06-17 14:08:28 -07:00
Nicholas Junge
cba5b13a39 Improve concretization errors for jnp indexing routines 2021-06-17 20:26:16 +02:00
Nicholas Junge
ccc8bb7f19 Add auxiliary data support to lax.custom_root 2021-06-17 19:25:46 +02:00
Jake VanderPlas
623d3c7ad3 [sparse] create experimental sparse submodule 2021-06-17 08:09:31 -07:00
jax authors
409684f9c2 Merge pull request #6982 from jakevdp:bcoo-dot-general 
PiperOrigin-RevId: 379810039
 2021-06-16 14:21:24 -07:00
jax authors
4bae894862 Merge pull request #6918 from colemanliyah:LRU_system 
PiperOrigin-RevId: 379736780
 2021-06-16 09:22:34 -07:00
jax authors
ea1be21ff1 Merge pull request #6927 from apaszke:xmap-loops 
PiperOrigin-RevId: 379736111
 2021-06-16 09:19:07 -07:00
Jake VanderPlas
6003f79063 [sparse] add bcoo_rdot_general 2021-06-16 08:54:30 -07:00
Liyah Coleman
0e9f7de6c8 Add LRU eviction policy to persisent compilation cache 2021-06-16 15:38:28 +00:00
Skye Wanderman-Milne
63355a9a25 Add jaxlib version check on newly-added NumpyLinalgTest.testEighZeroDiagonal test. 
This is testing behavior recently added to jaxlib, so fails on older jaxlibs.
 2021-06-15 20:44:52 +00:00
Adam Paszke
1e288b2f42 Add support for anonymous serial loops 2021-06-15 18:15:44 +00:00
jax authors
7f0abe409c Merge pull request #6967 from jakevdp:bcoo-matmul 
PiperOrigin-RevId: 379453035
 2021-06-15 03:28:41 -07:00
Jake VanderPlas
510a77700c Generalize BCOO.__matmul__ 2021-06-14 15:45:47 -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
Jake VanderPlas
2113d9c34d add bcoo_reduce_sum() function 2021-06-11 13:19:54 -07:00
jax authors
31e9c65f2a Merge pull request #6952 from jakevdp:hstack-reshape 
PiperOrigin-RevId: 378916971
 2021-06-11 11:43:38 -07:00
jax authors
8fcdb85f09 Merge pull request #6940 from jakevdp:fix-sinc 
PiperOrigin-RevId: 378904804
 2021-06-11 10:47:57 -07:00
Jake VanderPlas
17710c0711 add efficient path for array input to jnp.stack, jnp.[hvd]stack, jnp.column_stack 2021-06-11 10:42:06 -07:00
jax authors
3550732a74 Merge pull request #6946 from jakevdp:concat-reshape 
PiperOrigin-RevId: 378900980
 2021-06-11 10:32:58 -07:00
Peter Hawkins
b130257ee1 Drop support for NumPy 1.16. 2021-06-11 09:03:09 -04:00
Jake VanderPlas
0470f4f368 jnp.concatenate: add fast path based on lax.reshape for array inputs 2021-06-10 13:25:33 -07:00
Jake VanderPlas
80d8f2d56c jnp.sinc: fix NaNs at x=0 2021-06-10 09:14:07 -07:00
jax authors
69f6d5e3d2 Merge pull request #6781 from lukepfister:resize_weight 
PiperOrigin-RevId: 378550280
 2021-06-09 18:45:49 -07:00
jax authors
7d0bda604a Merge pull request #6930 from jakevdp:union1d-size 
PiperOrigin-RevId: 378530731
 2021-06-09 16:51:19 -07:00
Peter Hawkins
73c47dce6e [XLA] Revert to using the textbook algorithm to construct the 2x2 Jacobi rotations in Eigendecomposition. 
The current version is causing wrong outputs when the diagonal elements are exactly zero.

https://github.com/tensorflow/tensorflow/issues/50017

PiperOrigin-RevId: 378506479
 2021-06-09 14:56:45 -07:00
Jake VanderPlas
79d0852145 Add optional size argument to jnp.union1d for JIT compatibility 2021-06-09 11:36:34 -07:00
jax authors
2a1936e6f9 Merge pull request #6824 from jakevdp:sparse-bcoo 
PiperOrigin-RevId: 378437750
 2021-06-09 10:25:28 -07:00
Jake VanderPlas
69f29a631a Add experimental batched COO sparse format. 
This is an implementation of a batch-friendly multidimensional COO format for JAX. It contains implementations of four primitives (bcoo_todense, bcoo_fromdense, bcoo_extract, bcoo_dot_general), as well as batching, JVP, and transpose rules for each.

For convenience, this also adds class BCOO, which is a pytree wrapper around these.
 2021-06-09 09:10:53 -07:00
Jake VanderPlas
0f4f4102ce Add more complete test for jnp.nonzero size argument 2021-06-08 16:40:53 -07:00
Jake VanderPlas
f97e2f945f jnp.argwhere: add optional size parameter for JIT compatibility 2021-06-08 16:17:37 -07:00
Jake VanderPlas
022464b91b jnp.where: add optional size argument 2021-06-08 15:53:12 -07:00
Jake VanderPlas
1296dc3f1e jnp.flatnonzero: add optional size argument for JIT compatibility 2021-06-08 13:16:51 -07:00
Jake VanderPlas
d198ad0ac1 jnp.unique: add optional size argument for JIT compatibility 2021-06-08 11:31:42 -07:00