def test_one_dim_tensor_below_diag(self):
t = TensorBase(np.array([1, 2, 3, 4]))
self.assertTrue(syft.equal(syft.diag(t, -1), TensorBase([[0, 0, 0, 0, 0],
[1, 0, 0, 0, 0],
[0, 2, 0, 0, 0],
[0, 0, 3, 0, 0],
[0, 0, 0, 4, 0]])))
def test_two_dim_tensor_below_diag(self):
t = TensorBase(np.array([[0, 1], [2, 3]]))
self.assertTrue(syft.equal(syft.diag(t, -1), TensorBase([2])))
def test_two_dim_tensor_upper_diag(self):
t = TensorBase(np.array([[0, 1], [2, 3]]))
self.assertTrue(syft.equal(syft.diag(t, 1), TensorBase([1])))
def test_two_dim_tensor_main_diag(self):
t = TensorBase(np.array([[0, 1], [2, 3]]))
self.assertTrue(syft.equal(syft.diag(t, 0), TensorBase([0, 3])))
