mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-20 13:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
8,313 Commits 506 Branches 263 Tags
Commit Graph

5153 Commits

Author SHA1 Message Date
jax authors
eaaaa1106d Merge pull request #7077 from skye:version 
PiperOrigin-RevId: 381088906
 2021-06-23 12:20:31 -07:00
Skye Wanderman-Milne
444ee5e840 Update jax version to 0.2.15 2021-06-23 11:55:40 -07:00
Jake VanderPlas
c358da4e49 [sparse] add bcoo_transpose primitive 2021-06-23 10:04:27 -07:00
jax authors
28977761d5 Merge pull request #6849 from tomhennigan:changelist/376000598 
PiperOrigin-RevId: 381010658
 2021-06-23 05:46:01 -07:00
jax authors
b85e3a0107 Merge pull request #7049 from jakevdp:validate 
PiperOrigin-RevId: 380974236
 2021-06-23 00:44:56 -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
3b84f85700 Merge pull request #7058 from jakevdp:bcoo-serialization 
PiperOrigin-RevId: 380902758
 2021-06-22 15:13:31 -07:00
Jake VanderPlas
98ce6c5892 [sparse] cleanup validation function 2021-06-22 14:58:59 -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
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
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
Qiao Zhang
e8e6138c75 Add vmap rule for lax.clamp. 2021-06-21 14:59:34 -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
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
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
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
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
George Necula
8b320eca82 Test 6 for rename master to main 2021-06-17 16:33:06 +03:00
George Necula
26499a38e6 Merge pull request #7006 from gnecula:jax_test_5 
PiperOrigin-RevId: 379942050
 2021-06-17 15:25:38 +02:00
George Necula
4fa1506825 One more change, to tickle copybara 2021-06-17 16:08:54 +03:00
George Necula
268a242c65 Test 5 for rename master to main 2021-06-17 15:58:51 +03:00
George Necula
0ec53eec72 Test PR merged to main, before making main default 
Want to see the GitHub actions working, and copybara import
 2021-06-17 14:48:21 +03:00
George Necula
0d16f3cafd [jax2tf] Remove workaround for tf.math.atan2 not working on [b]float16. 
PiperOrigin-RevId: 379919580
 2021-06-17 03:21:51 -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
cac09d44c9 Merge pull request #6990 from zhangqiaorjc:fix_typos 
PiperOrigin-RevId: 379807021
 2021-06-16 14:07:27 -07:00
Qiao Zhang
57234e7eaa Fix typos and indent. 2021-06-16 11:10:42 -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