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

8265 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
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
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
Jake VanderPlas
623d3c7ad3 [sparse] create experimental sparse submodule 2021-06-17 08:09:31 -07:00
jax authors
6d805cae6e Merge pull request #7007 from gnecula:jax_change_6 
PiperOrigin-RevId: 379947590
 2021-06-17 06:51:49 -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
jax authors
97a5719fcb Merge pull request #7002 from gnecula:jax_test_3 
PiperOrigin-RevId: 379932113
 2021-06-17 04:56:03 -07: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
jax authors
daef97b1a1 Merge pull request #6999 from gnecula:rename_master 
PiperOrigin-RevId: 379929701
 2021-06-17 04:38:00 -07:00
George Necula
a6c16119c9 Renaming master to main: step 1. 
Here we just enable the GitHub actions to run on both master and main.
 2021-06-17 14:25:01 +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
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
jax authors
8f5a78404b Merge pull request #6964 from gnecula:tf_gather 
PiperOrigin-RevId: 379693316
 2021-06-16 04:42:51 -07:00
jax authors
ea711b31a2 Merge pull request #6986 from gnecula:call_tf_x64 
PiperOrigin-RevId: 379693059
 2021-06-16 04:39:10 -07:00
George Necula
d394f6917b Added clipping of OOB indices to match XLA behavior 2021-06-16 14:24:28 +03:00
George Necula
c956e83df0 [jax2tf] Add a best-effort conversion for gather with enable_xla=False 
In an attempt to support lax.gather in absence of XLA, e.g., for
TFjs converter, we add an alternative conversion for the case
when enable_xla=False. Previously, the conversion would fail in this case.

The important parts of the code are lifeted from an old PR #3486, which was
then removed in #4030 (to avoid having to think about out-of-bounds accesses).

There is still a discrepancy between the JAX and TF behavior for out-of-bounds
accesses.
 2021-06-16 14:24:28 +03:00
George Necula
e405a73663 [call_tf] Fixed the handling of x64 values in JAX_ENABLE_X64=False 
TensorFlow interprets Python scalars in x64 mode, or it may contain
explicit x64 computation. When used with call_tf, we canonicalize
inputs and outputs using the JAX rules.
 2021-06-16 13:25:56 +03:00
jax authors
ed96cac457 Merge pull request #6985 from gnecula:call_tf_grad 
PiperOrigin-RevId: 379681210
 2021-06-16 03:05:00 -07:00
George Necula
9831371e38 [jax2tf] Fix the custom gradients for call_tf 
The previous implementation did not work correctly when
the call_tf function was embedded in a larger JAX function
for which we take the gradient.

Also fixed handling of non-float args/results.
 2021-06-16 12:04:32 +03:00
jax authors
a82f0d8355 Merge pull request #6984 from jakevdp:bcoo-traced-repr 
PiperOrigin-RevId: 379592043
 2021-06-15 15:49:13 -07:00
Jake VanderPlas
591507a3eb [sparse] update sparse object repr to show tracing 2021-06-15 15:16:58 -07:00
jax authors
096f77dc70 Merge pull request #6983 from skye:test_fix 
PiperOrigin-RevId: 379572178
 2021-06-15 14:20:57 -07: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
fec82c8aa1 Merge pull request #6980 from jakevdp:fix-faq 
PiperOrigin-RevId: 379511290
 2021-06-15 09:53:12 -07:00