If you have problem with flatten you could try to see:

  

      github.com
  

  
    jonasrauber/eagerpy/blob/master/eagerpy/tensor/tensor.py#L609



      
      
              *args: Any,
              **kwargs: Any,
          ) -> Tuple[TensorType, Any, TensorType]:
              return self._value_and_grad_fn(f, has_aux=True)(self, *args, **kwargs)
          
          def __iter__(self: TensorType) -> Iterator[TensorType]:
              for i in range(len(self)):
                  yield self[i]
          
          @final
          def flatten(self: TensorType, start: int = 0, end: int = -1) -> TensorType:
              start = start % self.ndim
              end = end % self.ndim
              shape = self.shape[:start] + (-1,) + self.shape[end + 1 :]
              return self.reshape(shape)
          
          def __matmul__(self: TensorType, other: TensorType) -> TensorType:
              if self.ndim != 2 or other.ndim != 2:
                  raise ValueError(
                      f"matmul requires both tensors to be 2D, got {self.ndim}D and {other.ndim}D"
                  )