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

8313 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
jax authors
eaaaa1106d Merge pull request #7077 from skye:version 
PiperOrigin-RevId: 381088906
jax-v0.2.15 		2021-06-23 12:20:31 -07:00
jax authors
35557794fb Merge pull request #7070 from jakevdp:bcoo-transpose 
PiperOrigin-RevId: 381087809
 2021-06-23 12:15:34 -07:00
Skye Wanderman-Milne
444ee5e840 Update jax version to 0.2.15 2021-06-23 11:55:40 -07:00
jax authors
a251eccccb Merge pull request #7076 from jakevdp:gpu-install 
PiperOrigin-RevId: 381082192
 2021-06-23 11:49:08 -07:00
Jake VanderPlas
dfcaf0feb8 Improve support for pip install jax[cuda111] 2021-06-23 11:42:04 -07:00
jax authors
0934161aea Merge pull request #7071 from colemanliyah:fix_get_cache_key 
PiperOrigin-RevId: 381067436
 2021-06-23 10:48:44 -07:00
colemanliyah
d460a53a56 changed 1 to 0 2021-06-23 17:40:23 +00:00
Liyah Coleman
92a290bd56 fixed failing tests 2021-06-23 17:37:27 +00:00
jax authors
d1a25503a4 Merge pull request #7062 from skye:jaxlib_tpu 
PiperOrigin-RevId: 381063068
 2021-06-23 10:35:30 -07:00
jax authors
a68754103f Merge pull request #7072 from zhangqiaorjc:npv2 
PiperOrigin-RevId: 381062738
 2021-06-23 10:31:57 -07:00
Qiao Zhang
7c54b44dba Bump numpy versions for macos build. 2021-06-23 10:23:43 -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
a12a229546 Add pip install jax[tpu] configuration. 
This can be used on Cloud TPU VMs to automatically install compatible
versions of jax, jaxlib, and libtpu (the low-level library JAX uses to
access the TPU on Cloud TPU VMs).

The new install command requires a new jax release (`>=0.2.15`) and
jaxlib release (`>=0.1.68`) to work, since it requires both
cdfbd9dde1
and
ce2bc24996
to pick up the pip-installed libtpu. I'll update the README and Cloud
TPU VM documentation once these releases are out.
 2021-06-23 02:26:27 +00:00
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
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
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