Repository: OpenMined/PySyft
Pull request #177 (2017-08-23 05:36:15)

:sparkles: Implemented unsqueeze functionality
1. `unsqueeze` function in `math.py`
2. `unsqueeze` function for `TensorBase`
3. `unsqueeze_` Function for `TensorBase`
4. Respective unit tests


Patch:
diff --git a/syft/__init__.py b/syft/__init__.py
index cacd7ba6fb..13a92d338a 100644
--- a/syft/__init__.py
+++ b/syft/__init__.py
@@ -4,6 +4,7 @@ from syft import nn
 from syft.tensor import equal, TensorBase
 from syft.math import cumprod, cumsum, ceil, dot, matmul, addmm, addcmul
 from syft.math import addcdiv, addmv, addbmm, baddbmm
+from syft.math import unsqueeze
 
 s = str(he)
 s += str(nn)
@@ -11,3 +12,4 @@ s += str(nn)
 s += str(equal) + str(TensorBase) + str(cumprod) + str(cumsum) + str(ceil)
 s += str(dot) + str(matmul) + str(addmm) + str(addcmul) + str(addcdiv)
 s += str(addmv) + str(addbmm) + str(baddbmm)
+s += str(unsqueeze)
diff --git a/syft/math.py b/syft/math.py
index 263e3939af..0f8cf7c93c 100644
--- a/syft/math.py
+++ b/syft/math.py
@@ -7,9 +7,8 @@ from .tensor import TensorBase
 from .tensor import _ensure_tensorbase
 
 __all__ = [
-
     'cumprod', 'cumsum', 'ceil', 'dot', 'floor', 'matmul', 'addmm', 'addcmul',
-    'addcdiv', 'addmv', 'addbmm', 'baddbmm', 'sigmoid',
+    'addcdiv', 'addmv', 'addbmm', 'baddbmm', 'sigmoid', 'unsqueeze'
 ]
 
 
@@ -263,3 +262,19 @@ def baddbmm(tensor1, tensor2, mat, beta=1, alpha=1):
         mmul = np.matmul(tensor1.data, tensor2.data)
         out = (mat.data * beta) + (mmul * alpha)
         return TensorBase(out)
+
+
+def unsqueeze(tensor1, dim):
+    """
+    Performs 'unsqueeze' operation, returning a new tensor with a dimension
+    of size one inserted at the specified position.
+    """
+    tensor1 = _ensure_tensorbase(tensor1)
+    num_dims = len(tensor1.data.shape)
+
+    if dim >= num_dims or dim < 0:
+        print("dimension out of range")
+    elif tensor1.encrypted:
+        raise NotImplemented
+    else:
+        return TensorBase(np.expand_dims(tensor1.data, dim))
diff --git a/syft/tensor.py b/syft/tensor.py
index f6cb19e76a..227ab40028 100644
--- a/syft/tensor.py
+++ b/syft/tensor.py
@@ -418,6 +418,27 @@ class TensorBase(object):
             self.data += (mat.data * beta)
             return self
 
+    def unsqueeze(self, dim):
+        """
+        Returns expanded Tensor. An additional dimension of size one is added
+        to at index 'dim'.
+        """
+        return syft.unsqueeze(self.data, dim)
+
+    def unsqueeze_(self, dim):
+        """
+        Replaces with an expanded Tensor. An additional dimension of size one
+        is added to at index 'dim'.
+        """
+        num_dims = len(self.data.shape)
+
+        if dim >= num_dims or dim < 0:
+            print("dimension out of range")
+        elif self.encrypted:
+            raise NotImplemented
+        else:
+            self.data = np.expand_dims(self.data, dim)
+
     def exp(self):
         """Computes the exponential of each element in tensor."""
         if self.encrypted:


Test patch:
diff --git a/tests/test_math.py b/tests/test_math.py
index 8e153730e4..011bcfeb88 100644
--- a/tests/test_math.py
+++ b/tests/test_math.py
@@ -167,3 +167,13 @@ class baddbmmTests(unittest.TestCase):
         out = syft.baddbmm(t1, t2, mat, beta=2, alpha=2)
         self.assertTrue(np.array_equal(out.data, [[[62, 92], [96, 142]],
                                                   [[122, 184], [28, 42]]]))
+
+
+class unsqueezeTests(unittest.TestCase):
+    def testUnsqueeze(self):
+        t1 = TensorBase(np.arange(3 * 4 * 5).reshape((3, 4, 5)))
+        for i in range(len(t1.data.shape)):
+            out = syft.unsqueeze(t1, i)
+            expected_shape = list(t1.data.shape)
+            expected_shape.insert(i, 1)
+            self.assertTrue(np.array_equal(out.data.shape, expected_shape))
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index 195fb686c6..3c31305fb6 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -299,6 +299,25 @@ class baddbmmTests(unittest.TestCase):
                                                  [[122, 184], [28, 42]]]))
 
 
+class unsqueezeTests(unittest.TestCase):
+    def testUnsqueeze(self):
+        t1 = TensorBase(np.arange(3 * 4 * 5).reshape((3, 4, 5)))
+        for i in range(len(t1.data.shape)):
+            out = t1.unsqueeze(i)
+            expected_shape = list(t1.data.shape)
+            expected_shape.insert(i, 1)
+            self.assertTrue(np.array_equal(out.data.shape, expected_shape))
+
+    def testUnsqueeze_(self):
+        test_shape = (3, 4, 5)
+        for i in range(len(test_shape)):
+            t1 = TensorBase(np.arange(3 * 4 * 5).reshape(test_shape))
+            expected_shape = list(t1.data.shape)
+            expected_shape.insert(i, 1)
+            t1.unsqueeze_(i)
+            self.assertTrue(np.array_equal(t1.data.shape, expected_shape))
+
+
 class expTests(unittest.TestCase):
     def testexp(self):
         t3 = TensorBase(np.array([[[1, 3], [3, 5]],