Repository: openmined/pysyft
Issue: GammaTensor created incorrectly when calling .private()

The heuristic we use to check whether a tensor is a PhiTensor or a GammaTensor when calling `.private()` or `.annotate_with_dp_metadata` is whether len(data_subjects.sum()) > 1

This will raise false positives- GammaTensors have >1 data subject per value in the Tensor, but this check would also create a GammaTensor if data_subjects had 1 data subject per value, but they were all different.

Changed files:
packages/syft/src/syft/core/tensor/ancestors.py, packages/syft/tests/syft/core/tensor/private_test.py, packages/syft/src/syft/core/adp/data_subject_list.py

diff --git a/packages/syft/src/syft/core/adp/data_subject_list.py b/packages/syft/src/syft/core/adp/data_subject_list.py
index 9b39c1c9f..0d2343958 100644
--- a/packages/syft/src/syft/core/adp/data_subject_list.py
+++ b/packages/syft/src/syft/core/adp/data_subject_list.py
@@ -623,7 +623,7 @@ class DataSubjectArray:
         # per data point, we should make sure that we implement in such a way we expand
         # the datasubjects automatically for row to data point mapping.
         if not isinstance(input_subjects, np.ndarray):
-            input_subjects = np.array(input_subjects)
+            input_subjects = np.array(input_subjects, dtype=DataSubjectArray)
 
         data_map = (
             lambda x: DataSubjectArray([str(x)])
diff --git a/packages/syft/src/syft/core/tensor/ancestors.py b/packages/syft/src/syft/core/tensor/ancestors.py
index f11ec69e8..8c3365534 100644
--- a/packages/syft/src/syft/core/tensor/ancestors.py
+++ b/packages/syft/src/syft/core/tensor/ancestors.py
@@ -435,10 +435,9 @@ class PhiTensorAncestor(TensorChainManager):
             min_val, max_val, target_shape=self.child.shape
         )
 
-        unique_data_subjects = len(data_subjects.sum())
-        if unique_data_subjects == 1:
+        if any(len(x.item()) > 1 for x in np.nditer(data_subjects, flags=["refs_ok"])):
             self.replace_abstraction_top(
-                tensor_type=_PhiTensor(),
+                tensor_type=_GammaTensor(),
                 child=self.child,
                 min_vals=min_vals,
                 max_vals=max_vals,
@@ -446,7 +445,7 @@ class PhiTensorAncestor(TensorChainManager):
             )  # type: ignore
         else:
             self.replace_abstraction_top(
-                tensor_type=_GammaTensor(),
+                tensor_type=_PhiTensor(),
                 child=self.child,
                 min_vals=min_vals,
                 max_vals=max_vals,
diff --git a/packages/syft/tests/syft/core/tensor/private_test.py b/packages/syft/tests/syft/core/tensor/private_test.py
index 85cd100ad..2d3ec6be8 100644
--- a/packages/syft/tests/syft/core/tensor/private_test.py
+++ b/packages/syft/tests/syft/core/tensor/private_test.py
@@ -4,6 +4,7 @@ import pytest
 
 # syft absolute
 import syft as sy
+from syft.core.adp.data_subject_list import DataSubjectArray
 from syft.core.tensor.autodp.gamma_tensor import GammaTensor as GT
 from syft.core.tensor.autodp.phi_tensor import PhiTensor as PT
 from syft.core.tensor.lazy_repeat_array import lazyrepeatarray as lra
@@ -163,11 +164,38 @@ def test_2d_array(tensor: Tensor, low: int, high: int) -> None:
     assert len(private.child.data_subjects.sum()) == 1
 
 
+def test_phi(tensor: Tensor, low: int, high: int) -> None:
+    data_subjects = np.random.choice(["Optimus Prime", "Bumblebee"], (5, 5))
+    # Make sure there's at least one of "Optimus Prime" and "Bumblebee" to prevent
+    # the 1/2^24 chance of failure
+    data_subjects[0, 0] = "Optimus Prime"
+    data_subjects[4, 4] = "Bumblebee"
+
+    private = tensor.private(
+        min_val=low,
+        max_val=high,
+        data_subjects=data_subjects,
+    )
+    assert isinstance(private, Tensor)
+    assert isinstance(private.child, PT)
+    assert isinstance(private.child.min_vals, lra)
+    assert isinstance(private.child.max_vals, lra)
+    assert private.child.min_vals.shape == private.child.shape
+    assert private.child.max_vals.shape == private.child.shape
+    assert isinstance(private.child.data_subjects, np.ndarray)
+    assert private.child.data_subjects.shape == private.child.shape
+    assert len(private.child.data_subjects.sum()) == 2
+
+
 def test_gamma(tensor: Tensor, low: int, high: int) -> None:
+    data_subjects = np.random.choice(["Optimus Prime", "Bumblebee"], (5, 5)).tolist()
+    data_subjects = [[DataSubjectArray([x]) for x in row] for row in data_subjects]
+    data_subjects[0][0] = DataSubjectArray(["Optimus Prime", "Bumblebee"])
+
     private = tensor.private(
         min_val=low,
         max_val=high,
-        data_subjects=np.random.choice(["Optimus Prime", "Bumblebee"], (5, 5)),
+        data_subjects=data_subjects,
     )
     assert isinstance(private, Tensor)
     assert isinstance(private.child, GT)