Repository: OpenMined/PySyft

Issue:
[SMPC] Implement matmul
## Feature Description
Implement matrix multiplication for the SMPC.
Follow the example of `mul` operation from the `mpc_tensor.py` file.

## Is your feature request related to a problem?
Yes. Have the possiblity to run `matrix multiplication`.

## What alternatives have you considered?

## Acceptance Criteria
- [ ] Implement the functionality
- [ ] Add tests
- [ ] Update the SMPC Notebook


Discussion:

Resolving patch:
diff --git a/packages/syft/src/syft/core/tensor/__init__.py b/packages/syft/src/syft/core/tensor/__init__.py
--- a/packages/syft/src/syft/core/tensor/__init__.py
+++ b/packages/syft/src/syft/core/tensor/__init__.py
@@ -178,6 +178,14 @@ def create_tensor_ast(client: Optional[AbstractNodeClient] = None) -> Globals:
             "syft.core.tensor.smpc.share_tensor.ShareTensor.__mul__",
             "syft.core.tensor.smpc.share_tensor.ShareTensor",
         ),
+        (
+            "syft.core.tensor.smpc.share_tensor.ShareTensor.__matmul__",
+            "syft.core.tensor.smpc.share_tensor.ShareTensor",
+        ),
+        (
+            "syft.core.tensor.smpc.share_tensor.ShareTensor.__rmatmul__",
+            "syft.core.tensor.smpc.share_tensor.ShareTensor",
+        ),
         (
             "syft.core.tensor.smpc.share_tensor.ShareTensor.sum",
             "syft.core.tensor.smpc.share_tensor.ShareTensor",
diff --git a/packages/syft/src/syft/core/tensor/smpc/mpc_tensor.py b/packages/syft/src/syft/core/tensor/smpc/mpc_tensor.py
--- a/packages/syft/src/syft/core/tensor/smpc/mpc_tensor.py
+++ b/packages/syft/src/syft/core/tensor/smpc/mpc_tensor.py
@@ -364,8 +364,11 @@ def __apply_private_op(self, other: MPCTensor, op_str: str) -> List[ShareTensor]
 
     def __apply_public_op(self, y: Any, op_str: str) -> List[ShareTensor]:
         op = getattr(operator, op_str)
-        if op_str in {"mul", "matmul", "add", "sub"}:
+        if op_str in {"mul", "matmul"}:
             res_shares = [op(share, y) for share in self.child]
+        elif op_str in {"add", "sub"}:
+            res_shares = self.child
+            res_shares[0] = op(res_shares[0], y)
         else:
             raise ValueError(f"{op_str} not supported")
 
@@ -481,6 +484,24 @@ def mul(
 
         return res
 
+    def matmul(
+        self, y: Union[int, float, np.ndarray, torch.tensor, "MPCTensor"]
+    ) -> MPCTensor:
+        """Apply the "matmul" operation between "self" and "y"
+
+        Args:
+            y (Union[int, float, np.ndarray, torch.tensor, "MPCTensor"]): self @ y
+
+        Returns:
+            MPCTensor: Result of the opeartion.
+        """
+        if isinstance(y, ShareTensor):
+            raise ValueError("Private matmul not supported yet")
+
+        res = self.__apply_op(y, "matmul")
+
+        return res
+
     def __str__(self) -> str:
         res = "MPCTensor"
         for share in self.child:
@@ -529,6 +550,7 @@ def put(
     __rsub__ = rsub
     __mul__ = mul
     __rmul__ = mul
+    __matmul__ = matmul
 
 
 @implements(MPCTensor, np.add)
diff --git a/packages/syft/src/syft/core/tensor/smpc/share_tensor.py b/packages/syft/src/syft/core/tensor/smpc/share_tensor.py
--- a/packages/syft/src/syft/core/tensor/smpc/share_tensor.py
+++ b/packages/syft/src/syft/core/tensor/smpc/share_tensor.py
@@ -315,7 +315,10 @@ def matmul(
         Returns:
             ShareTensor: Result of the operation.
         """
-        ShareTensor.sanity_checks(y)
+        if isinstance(y, ShareTensor):
+            raise ValueError("Private matmul not supported yet")
+
+        ShareTensor.sanity_check(y)
         new_share = self.apply_function(y, "matmul")
         return new_share
 
@@ -328,8 +331,12 @@ def rmatmul(self, y: torch.Tensor) -> "ShareTensor":
         Returns:
             ShareTensor. Result of the operation.
         """
-        ShareTensor.sanity_checks(y)
-        return y.matmul(self)
+        if isinstance(y, ShareTensor):
+            raise ValueError("Private matmul not supported yet")
+
+        ShareTensor.sanity_check(y)
+        new_share = y.apply_function(self, "matmul")
+        return new_share
 
     # TRASK: commenting out because ShareTEnsor doesn't appear to have .session_uuid or .config
     # def div(


Test patch:
diff --git a/packages/syft/tests/syft/core/tensor/smpc/smpc_tensor_test.py b/packages/syft/tests/syft/core/tensor/smpc/smpc_tensor_test.py
--- a/packages/syft/tests/syft/core/tensor/smpc/smpc_tensor_test.py
+++ b/packages/syft/tests/syft/core/tensor/smpc/smpc_tensor_test.py
@@ -90,6 +90,25 @@ def test_mpc_private_public_op(op_str: str, public_value_type: str) -> None:
     assert (res == expected).all()
 
 
+@pytest.mark.parametrize("op_str", ["matmul"])
+def test_mpc_matmul_op(op_str: str) -> None:
+    value_1 = np.array([[1, 7], [3, -7]], dtype=np.int32)
+    value_2 = np.array([[6, 2], [-6, 5]], dtype=np.int32)
+
+    remote_value_1 = clients[0].syft.core.tensor.tensor.Tensor(value_1)
+
+    mpc_tensor_1 = MPCTensor(
+        parties=clients, secret=remote_value_1, shape=(2, 2), seed_shares=52
+    )
+
+    op = getattr(operator, op_str)
+
+    res = op(mpc_tensor_1, value_2).reconstruct()
+    expected = op(value_1, value_2)
+
+    assert (res == expected).all()
+
+
 @pytest.mark.parametrize("public_value_type", ["int", "torch_tensor", "numpy_array"])
 @pytest.mark.parametrize("op_str", ["add", "sub", "mul"])
 def test_mpc_public_private_op(op_str: str, public_value_type: str) -> None: