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

8296 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
jax authors
cdfbd9dde1 Merge pull request #7061 from skye:libtpu_init 
PiperOrigin-RevId: 380932769
jaxlib-v0.1.68 		2021-06-22 18:11:45 -07:00
Skye Wanderman-Milne
ba972f0207 On Cloud TPU, use pip-installed libtpu instead of system default if applicable. 2021-06-22 23:51:57 +00:00
jax authors
6ae2f6d5a4 Merge pull request #7060 from zhangqiaorjc:npv 
PiperOrigin-RevId: 380918414
 2021-06-22 16:37:50 -07:00
Qiao Zhang
f3c8a22d66 Bump np versions for build script. 2021-06-22 16:29:56 -07:00
jax authors
3b84f85700 Merge pull request #7058 from jakevdp:bcoo-serialization 
PiperOrigin-RevId: 380902758
 2021-06-22 15:13:31 -07:00
jax authors
afe2d908ce Merge pull request #7042 from zhangqiaorjc:tfrt_cpu_default 
PiperOrigin-RevId: 380899063
 2021-06-22 14:57:49 -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
Qiao Zhang
132a54228b Add flag to select tfrt backend for CPU. 2021-06-22 14:12:36 -07:00
jax authors
0ed1b2b90e Merge pull request #7045 from zhangqiaorjc:win_build 
PiperOrigin-RevId: 380884480
 2021-06-22 13:48:47 -07:00
Qiao Zhang
18c7610e96 Catch CalledProcessError in build script. 2021-06-22 13:36:14 -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
jax authors
4b8d948da7 Merge pull request #7056 from zhangqiaorjc:updatexla 
PiperOrigin-RevId: 380877984
 2021-06-22 13:17:08 -07:00
Qiao Zhang
bfc2ae2161 Update XLA. 2021-06-22 13:07:24 -07:00
jax authors
23b6b0bbb4 Merge pull request #7051 from jakevdp:pin-numpy 
PiperOrigin-RevId: 380871374
 2021-06-22 12:48:08 -07:00
Jake VanderPlas
0c91be7b46 CI: temporarily pin numpy to <1.21 2021-06-22 11:15:16 -07: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
jax authors
8e14d53bdb Merge pull request #7041 from lamblin:patch-1 
PiperOrigin-RevId: 380787775
 2021-06-22 05:39:30 -07:00
Marc van Zee
1bba5d7f10 Implements dynamic_update_slice when enable_xla=False and adds tests. 
PiperOrigin-RevId: 380781017
 2021-06-22 04:43:48 -07:00
Marc van Zee
e50276019d Implements more slicing support for lax.gather in jax2tf when enable_xla=False, and adds tests. 
PiperOrigin-RevId: 380753589
 2021-06-22 01:07:03 -07:00
Liyah Coleman
a87a291afd Implemented get_cache_key funtion with unit tests 2021-06-21 22:49:09 +00:00
jax authors
6382cee708 Merge pull request #7034 from zhangqiaorjc:vmapped_clamp 
PiperOrigin-RevId: 380677496
 2021-06-21 15:27:09 -07:00
Qiao Zhang
e8e6138c75 Add vmap rule for lax.clamp. 2021-06-21 14:59:34 -07:00
jax authors
3787b56c15 Merge pull request #7040 from jakevdp:fix-random-test 
PiperOrigin-RevId: 380658330
 2021-06-21 13:59:14 -07:00
Jake VanderPlas
71bf410055 Fix random test for scipy>=1.17.0 2021-06-21 13:44:54 -07:00
jax authors
3f768ee27a Merge pull request #7036 from lesteve:patch-1 
PiperOrigin-RevId: 380639325
 2021-06-21 12:33:18 -07:00
Pascal Lamblin
963d4c0007
Fix typo in link 2021-06-21 14:48:33 -04:00
jax authors
8bb6bdc4b8 Merge pull request #7025 from jakevdp:bcoo-batch 
PiperOrigin-RevId: 380561985
 2021-06-21 06:14:00 -07:00
George Necula
64974f309a [host_callback] Make logging statements conditional. 
Misuse of conditional logging (logging.vlog) led to significant time wasted
even if logging was off.

PiperOrigin-RevId: 380559720
 2021-06-21 05:56:23 -07:00
Loïc Estève
6192516dfd
Use scikit-learn rather than sklearn in docs/requirements.txt 2021-06-21 12:03:44 +02:00
jax authors
50f48e8a48 Merge pull request #7027 from jakevdp:transpose-validation 
PiperOrigin-RevId: 380364559
 2021-06-19 09:38:21 -07:00
Jake VanderPlas
7d2057b5c7 BUG: fix validation of permutation in transpose 2021-06-18 21:54:28 -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
jax authors
9450b8f8f9 Merge pull request #6954 from lukepfister:complex_convolve 
PiperOrigin-RevId: 380260457
 2021-06-18 14:05:58 -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
jax authors
8e0ec1124e Merge pull request #6996 from jakevdp:sparse-grad 
PiperOrigin-RevId: 380201282
 2021-06-18 09:26:59 -07:00
jax authors
48c046c170 Merge pull request #7012 from jakevdp:fix-random-bits 
PiperOrigin-RevId: 380194114
 2021-06-18 08:47:50 -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
jax authors
9f8fb243f9 Merge pull request #7014 from gnecula:master_main 
PiperOrigin-RevId: 380144198
 2021-06-18 02:27:50 -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
George Necula
4e58106fcb [jax2tf] Add back the limitation for dot_general with preferred_element_type on GPU 
That case is still broken, it fails non-deterministically on GPU

PiperOrigin-RevId: 380122429
 2021-06-17 23:23:13 -07:00
Jake VanderPlas
d20ecf5db8 random: generate large sequences via vmap rather than Python loops 2021-06-17 16:10:30 -07:00
jax authors
1d0c461ea6 Merge pull request #6995 from jakevdp:sparse-submodule 
PiperOrigin-RevId: 380040999
 2021-06-17 14:08:28 -07:00
jax authors
27ee4b64b0 Merge pull request #7008 from njunge94:index-concretizations 
PiperOrigin-RevId: 380024740
 2021-06-17 12:55:25 -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