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

8325 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Skye Wanderman-Milne
3da8a4cd86 Update jax version to 0.2.16 2021-06-23 14:45:45 -07:00
jax authors
c985d76ee5 Merge pull request #7082 from zhangqiaorjc:gtsv2fix 
PiperOrigin-RevId: 381118896
 2021-06-23 14:44:41 -07:00
jax authors
4c1856e1b9 Merge pull request #7083 from skye:fix_install 
PiperOrigin-RevId: 381117784
 2021-06-23 14:39:04 -07:00
Qiao Zhang
7cc277d634 Test for gtsv2 attr on cusparse. 2021-06-23 14:22:08 -07:00
Skye Wanderman-Milne
55276d15e4 Fix pip install jax[tpu] 
* Updates jax_releases.html index to include libtpu wheels
* Change [tpu] extras to specify `libtpu-nightly` instead of wheel URL

The full install command will now be:
`pip install pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_releases.html`
(similar to the cuda install commands)

I've already pushed an updated jax_releases.html to the jax-releases GCS bucket.
 2021-06-23 14:13:15 -07:00
Skye Wanderman-Milne
2460f91561 Revert "Improve support for pip install jax[cuda111]" 
This reverts commit dfcaf0feb81db0d2594b9a78154e8679b0366533.
 2021-06-23 14:08:51 -07:00
jax authors
9493b315bc Merge pull request #7080 from jakevdp:test-dynamic-slice 
PiperOrigin-RevId: 381109882
 2021-06-23 14:00:54 -07:00
jax authors
5bfcb51f59 Merge pull request #7059 from zhangqiaorjc:cpp_fix 
PiperOrigin-RevId: 381104915
 2021-06-23 13:36:56 -07:00
Jake VanderPlas
27fc797a67 Add dynamic slice U8 index test 2021-06-23 13:29:15 -07:00
jax authors
92e8ec426c Merge pull request #7078 from zhangqiaorjc:bjaxlib 
PiperOrigin-RevId: 381099458
 2021-06-23 13:09:10 -07:00
Qiao Zhang
a22841b6bb Bump jaxlib ver to 0.1.68. 2021-06-23 12:37:56 -07:00
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
Qiao Zhang
cf64f840e1 Remove redundant dep in jax cpp example target. 2021-06-22 14:49:19 -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