Repository: OpenMined/PySyft
Pull request #374 (2017-10-24 19:45:20)

changes multinomial to match PyTorch's API
<!--
Thanks for submitting a PR. Please make sure  that you have read the contribution guidelines: https://github.com/OpenMined/Docs/blob/master/contributing/guidelines.md
-->
#### Reference Issue

<!-- Example: Fixes #334 -->


#### What does your implementation fix? 
Addresses the issue discussed in PR #367 comments

#### Explain your changes.
It changes multinomial implementation to match PyTorch's API

#### include other comments?(if any)


<!--
Please note that we are a small team of volunteers so patience is
essential; We value all the contributions, no matter how small they are. 
If we are slow to review, either the pull request needs 
some tweaking,convincing, etc. 
or probably the reviewers could be busy. In any
case, we hope for your understanding in the process.

However If you feel that  this needs to be fastened up.
Reachout to us on our slack #pysyft channel.

Thanks for the contribution!
-->


Patch:
diff --git a/syft/math.py b/syft/math.py
index bdbc3e1223..dd20a5e0c4 100644
--- a/syft/math.py
+++ b/syft/math.py
@@ -845,3 +845,37 @@ def renorm(tensor1, p, dim, maxnorm):
         scalar_reshaped = scalar.reshape(dim_array)
         out = tensor1 * scalar_reshaped
     return TensorBase(out)
+
+
+def multinomial(tensor, num_samples, replacement=False):
+    """
+    Returns Tensor with random numbers from the Multinomial Distribution.
+
+    Returns Tensor with random numbers
+    from a multinomial distribution with probability
+    specified by ``input``(arr_like), number of draws specified by ``num_samples``,
+    and whether to replace the draws specified by replacement.
+
+    The ``input`` Tensor should be a tensor containing probabilities to
+    be used for drawing the multinomial random number.
+    The values of ``input`` do not need to sum to one (in which case we use the values as weights),
+    but must be non-negative and have a non-zero sum.
+
+    Parameters
+    ----------
+    tensor: TensorBase
+        Weights for the multinomial distribution
+    num_samples: Int
+        Number of samples to be drawn. If replacement is false, this must be lower than the length of p.
+    replacement: bool, optional
+        Whether to draw with replacement or not
+
+    Returns
+    -------
+    Output Tensor
+    """
+    if tensor.encrypted:
+        return NotImplemented
+    p = _ensure_tensorbase(tensor)
+    p = p / p.sum()
+    return TensorBase(np.random.choice(len(p), num_samples, replacement, p.data))
diff --git a/syft/tensor.py b/syft/tensor.py
index 3f3288a7f0..67227024be 100644
--- a/syft/tensor.py
+++ b/syft/tensor.py
@@ -1930,24 +1930,22 @@ class TensorBase(object):
         self.data = np.random.binomial(1, p.data)
         return self
 
-    def multinomial(self, p, num_samples, replacement=False):
+    def multinomial(self, num_samples, replacement=False):
         """
         Returns Tensor with random numbers from the Multinomial Distribution.
 
         Returns Tensor with random numbers
         from a multinomial distribution with probability
-        specified by p(arr_like), number of draws specified by num_samples,
+        specified by ``self``, number of draws specified by num_samples,
         and whether to replace the draws specified by replacement.
 
-        The p Tensor should be a tensor containing probabilities to
+        The ``self`` should be a tensor containing probabilities to
         be used for drawing the multinomial random number.
-        The values of p do not need to sum to one (in which case we use the values as weights),
+        The values of ``self`` do not need to sum to one (in which case we use the values as weights),
         but must be non-negative and have a non-zero sum.
 
         Parameters
         ----------
-        p: TensorBase
-            Weights for the multinomial distribution
         num_samples: Int
             Number of samples to be drawn. If replacement is false, this must be lower than the length of p.
         replacement: bool, optional
@@ -1957,11 +1955,7 @@ class TensorBase(object):
         -------
         Output Tensor
         """
-        if self.encrypted:
-            return NotImplemented
-        p = _ensure_tensorbase(p)
-        p = p / p.sum()
-        return TensorBase(np.random.choice(len(p), num_samples, replacement, p.data))
+        return syft.math.multinomial(self, num_samples=num_samples, replacement=replacement)
 
     def uniform_(self, low=0, high=1):
         """


Test patch:
diff --git a/tests/test_math.py b/tests/test_math.py
index 86167ffe10..a8445e7512 100644
--- a/tests/test_math.py
+++ b/tests/test_math.py
@@ -349,3 +349,11 @@ class RenormTests(unittest.TestCase):
         t2 = syft.math.renorm(t1, 1, 2, 8)
         self.assertTrue(np.allclose(t2, np.array([[[1.0, 1.230770, 1.333333], [4.0, 3.076923, 2.666667]],
                                                   [[2.0, 1.846154, 1.777778], [1.0, 1.846154, 2.222222]]])))
+
+
+class MultinomialTests(unittest.TestCase):
+    def test_multinomial(self):
+        t1 = TensorBase([0.25, .25, .25, .25])
+        t2 = syft.math.multinomial(t1, len(t1))
+        self.assertTupleEqual((len(t1),), t2.shape())
+        self.assertTrue(np.all(t2.data >= 0) and np.all(t2.data <= len(t1)))
diff --git a/tests/test_tensor.py b/tests/test_tensor.py
index 02b4f06d27..686b3798d3 100644
--- a/tests/test_tensor.py
+++ b/tests/test_tensor.py
@@ -732,11 +732,10 @@ class BernoulliTests(unittest.TestCase):
 
 class MultinomialTests(unittest.TestCase):
     def test_multinomial(self):
-        p = TensorBase([0.25, .25, .25, .25])
-        t1 = TensorBase(np.zeros((1, 5)))
-        t2 = t1.multinomial(p, len(p))
-        self.assertTupleEqual((len(p),), t2.shape())
-        self.assertTrue(np.all(t2.data >= 0) and np.all(t2.data <= len(p)))
+        t1 = TensorBase([0.25, .25, .25, .25])
+        t2 = t1.multinomial(len(t1))
+        self.assertTupleEqual((len(t1),), t2.shape())
+        self.assertTrue(np.all(t2.data >= 0) and np.all(t2.data <= len(t1)))
 
 
 class CauchyTests(unittest.TestCase):