Merge pull request #43 from chcodes/fix-scalars-to-cuda

Fixed scalars being wrong input to ElementwiseKernels
mikgroup · Apr 27, 2020 · 4033dea · 4033dea
2 parents cbbe549 + a659737
commit 4033dea
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/sigpy/thresh.py b/sigpy/thresh.py
@@ -26,10 +26,14 @@ def soft_thresh(lamda, input):
         array: soft-thresholded result.
 
     """
-    xp = backend.get_array_module(input)
+    device = backend.get_device(input)
+    xp = device.xp
     if xp == np:
         return _soft_thresh(lamda, input)
     else:  # pragma: no cover
+        if np.isscalar(lamda):
+            lamda = backend.to_device(lamda, device)
+
         return _soft_thresh_cuda(lamda, input)
 
 
@@ -44,10 +48,14 @@ def hard_thresh(lamda, input):
         array: hard-thresholded result.
 
     """
-    xp = backend.get_array_module(input)
+    device = backend.get_device(input)
+    xp = device.xp
     if xp == np:
         return _hard_thresh(lamda, input)
     else:  # pragma: no cover
+        if np.isscalar(lamda):
+            lamda = backend.to_device(lamda, device)
+
         return _hard_thresh_cuda(lamda, input)
 
 

diff --git a/sigpy/util.py b/sigpy/util.py
@@ -429,15 +429,19 @@ def axpy(y, a, x):
 
     Args:
         y (array): Output array.
-        a (scalar): Input scalar.
+        a (scalar or array): Input scalar.
         x (array): Input array.
 
     """
-    xp = backend.get_array_module(y)
+    device = backend.get_device(y)
+    xp = device.xp
 
     if xp == np:
         _axpy(y, a, x, out=y)
     else:
+        if np.isscalar(a):
+            a = backend.to_device(a, device)
+
         _axpy_cuda(a, x, y)
 
 
@@ -446,14 +450,18 @@ def xpay(y, a, x):
 
     Args:
         y (array): Output array.
-        a (scalar): Input scalar.
+        a (scalar or array): Input scalar.
         x (array): Input array.
     """
-    xp = backend.get_array_module(y)
+    device = backend.get_device(y)
+    xp = device.xp
 
     if xp == np:
         _xpay(y, a, x, out=y)
     else:
+        if np.isscalar(a):
+            a = backend.to_device(a, device)
+
         _xpay_cuda(a, x, y)
 
 
@@ -484,4 +492,4 @@ def _xpay(y, a, x):
         """
         y = x + (T) a * y;
         """,
-        name='axpy')
+        name='xpay')