Implement eagerpy.diag

eagerpy/framework.py

@@ -175,6 +175,10 @@ def transpose(t: TensorType, axes: Optional[Axes] = None) -> TensorType:
     return t.transpose(axes=axes)
 
 
+def diag(t: TensorType, k: int = 0) -> TensorType:
+    return t.diag(k)
+
+
 @overload
 def logical_and(x: TensorType, y: TensorOrScalar) -> TensorType:
     ...

eagerpy/tensor/jax.py

@@ -256,6 +256,9 @@ def transpose(self: TensorType, axes: Optional[Axes] = None) -> TensorType:
             axes = tuple(range(self.ndim - 1, -1, -1))
         return type(self)(np.transpose(self.raw, axes=axes))
 
+    def diag(self: TensorType, k: int = 0) -> TensorType:
+        return type(self)(np.diag(self.raw, k=k))
+
     def all(
         self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
     ) -> TensorType:

eagerpy/tensor/numpy.py

@@ -197,6 +197,9 @@ def transpose(self: TensorType, axes: Optional[Axes] = None) -> TensorType:
             axes = tuple(range(self.ndim - 1, -1, -1))
         return type(self)(np.transpose(self.raw, axes=axes))
 
+    def diag(self: TensorType, k: int = 0) -> TensorType:
+        return type(self)(np.diag(self.raw, k=k))
+
     def all(
         self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
     ) -> TensorType:

eagerpy/tensor/pytorch.py

@@ -281,6 +281,9 @@ def transpose(self: TensorType, axes: Optional[Axes] = None) -> TensorType:
             axes = tuple(range(self.ndim - 1, -1, -1))
         return type(self)(self.raw.permute(*axes))
 
+    def diag(self: TensorType, k: int = 0) -> TensorType:
+        return type(self)(torch.diag(self.raw, diagonal=k))
+
     def all(
         self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
     ) -> TensorType:

eagerpy/tensor/tensor.py

@@ -373,6 +373,10 @@ def _stack(
     def transpose(self: TensorType, axes: Optional[Axes] = None) -> TensorType:
         ...
 
+    @abstractmethod
+    def diag(self: TensorType, k: int = 0) -> TensorType:
+        ...
+
     @abstractmethod
     def take_along_axis(self: TensorType, index: TensorType, axis: int) -> TensorType:
         ...

eagerpy/tensor/tensorflow.py

@@ -263,6 +263,12 @@ def transpose(self: TensorType, axes: Optional[Axes] = None) -> TensorType:
             axes = tuple(range(self.ndim - 1, -1, -1))
         return type(self)(tf.transpose(self.raw, perm=axes))
 
+    def diag(self: TensorType, k: int = 0) -> TensorType:
+        if len(self.shape) == 1:
+            return type(self)(tf.linalg.diag(self.raw, k=k))
+        else:
+            return type(self)(tf.linalg.diag_part(self.raw, k=k))
+
     def all(
         self: TensorType, axis: Optional[AxisAxes] = None, keepdims: bool = False
     ) -> TensorType:

tests/test_main.py

@@ -239,6 +239,59 @@ def test_transpose_1d(dummy: Tensor) -> None:
     assert (ep.transpose(t) == t).all()
 
 
+def test_diag(dummy: Tensor) -> None:
+
+    t = ep.arange(dummy, 1, 5).float32()
+    d = ep.diag(t)
+    assert d.shape == (4, 4)
+    assert (d.flatten()[[0, 5, 10, 15]] == t).all()
+    assert (ep.index_update(d.flatten(), [0, 5, 10, 15], ep.zeros(dummy, 4)) == 0).all()
+    assert (ep.diag(d) == t).all()
+
+    d = ep.diag(t, k=1)
+    assert d.shape == (5, 5)
+    assert (d.flatten()[[1, 7, 13, 19]] == t).all()
+    assert (ep.index_update(d.flatten(), [1, 7, 13, 19], ep.zeros(dummy, 4)) == 0).all()
+    assert (ep.diag(d, k=1) == t).all()
+
+    d = ep.diag(t, k=2)
+    assert d.shape == (6, 6)
+    assert (d.flatten()[[2, 9, 16, 23]] == t).all()
+    assert (ep.index_update(d.flatten(), [2, 9, 16, 23], ep.zeros(dummy, 4)) == 0).all()
+    assert (ep.diag(d, k=2) == t).all()
+
+    d = ep.diag(t, k=-1)
+    assert d.shape == (5, 5)
+    assert (d.flatten()[[5, 11, 17, 23]] == t).all()
+    assert (
+        ep.index_update(d.flatten(), [5, 11, 17, 23], ep.zeros(dummy, 4)) == 0
+    ).all()
+    assert (ep.diag(d, k=-1) == t).all()
+
+    d = ep.diag(t, k=-2)
+    assert d.shape == (6, 6)
+    assert (d.flatten()[[12, 19, 26, 33]] == t).all()
+    assert (
+        ep.index_update(d.flatten(), [12, 19, 26, 33], ep.zeros(dummy, 4)) == 0
+    ).all()
+    assert (ep.diag(d, k=-2) == t).all()
+
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    assert (ep.diag(t) == t.flatten()[[0, 4, 8]]).all()
+
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    assert (ep.diag(t, k=1) == t.flatten()[[1, 5]]).all()
+
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    assert (ep.diag(t, k=2) == t.flatten()[[2]]).all()
+
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    assert (ep.diag(t, k=-1) == t.flatten()[[3, 7]]).all()
+
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    assert (ep.diag(t, k=-2) == t.flatten()[[6]]).all()
+
+
 def test_onehot_like_raises(dummy: Tensor) -> None:
     t = ep.arange(dummy, 18).float32().reshape((6, 3))
     indices = ep.arange(t, 6) // 2
@@ -1227,6 +1280,42 @@ def test_transpose(dummy: Tensor) -> Tensor:
     return ep.transpose(t)
 
 
+@compare_all
+def test_diag_1(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 4).float32()
+    return ep.diag(t)
+
+
+@compare_all
+def test_diag_2(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 4).float32()
+    return ep.diag(t, k=2)
+
+
+@compare_all
+def test_diag_3(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 4).float32()
+    return ep.diag(t, k=-2)
+
+
+@compare_all
+def test_diag_4(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    return ep.diag(t)
+
+
+@compare_all
+def test_diag_5(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    return ep.diag(t, k=2)
+
+
+@compare_all
+def test_diag_6(dummy: Tensor) -> Tensor:
+    t = ep.arange(dummy, 9).float32().reshape((3, 3))
+    return ep.diag(t, k=-2)
+
+
 @compare_all
 def test_transpose_axes(dummy: Tensor) -> Tensor:
     t = ep.arange(dummy, 60).float32().reshape((3, 4, 5))