code: def diagonal(tensor: Tensor) -> Tensor:
      assert all(
          sh == sh[0] for sh in tensor.shape
      ), "matrix should be square"
      assert tensor.ndim == 2, "only support 2 ndim tensors"
      return (Tensor.eye(tensor.shape[0]) * tensor).sum(0)
  
imports: import math, unittest, random, copy
  import numpy
  from tinygrad import Tensor, dtypes, Device
  from tinygrad.lazy import LazyBuffer
  from tinygrad.shape.shapetracker import ShapeTracker
  from tinygrad.shape.view import View
  from tinygrad.helpers import CI
label_desc: Use image node_id 8 for calling a global function with example usage: diagonal(tensor) and returns: unknown