code: def fill_diagonal(a, val, wrap=False):
      """Fill the main diagonal of the given tensor of any dimensionality.
  
      For a tensor `a` with ``a.ndim >= 2``, the diagonal is the list of
      locations with indices ``a[i, ..., i]`` all identical. This function
      modifies the input tensor in-place, it does not return a value.
  
      Parameters
      ----------
      a : Tensor, at least 2-D.
        Tensor whose diagonal is to be filled, it gets modified in-place.
  
      val : scalar
        Value to be written on the diagonal, its type must be compatible with
        that of the tensor a.
  
      wrap : bool
        For tall matrices in NumPy version up to 1.6.2, the
        diagonal "wrapped" after N columns. You can have this behavior
        with this option. This affects only tall matrices.
  
      See also
      --------
      diag_indices, diag_indices_from
  
      Notes
      -----
  
      This functionality can be obtained via `diag_indices`, but internally
      this version uses a much faster implementation that never constructs the
      indices and uses simple slicing.
  
      Examples
      --------
      >>> import mars.tensor as mt
      >>> a = mt.zeros((3, 3), int)
      >>> mt.fill_diagonal(a, 5)
      >>> a.execute()
      array([[5, 0, 0],
             [0, 5, 0],
             [0, 0, 5]])
  
      The same function can operate on a 4-D tensor:
  
      >>> a = mt.zeros((3, 3, 3, 3), int)
      >>> mt.fill_diagonal(a, 4)
  
      We only show a few blocks for clarity:
  
      >>> a[0, 0].execute()
      array([[4, 0, 0],
             [0, 0, 0],
             [0, 0, 0]])
      >>> a[1, 1].execute()
      array([[0, 0, 0],
             [0, 4, 0],
             [0, 0, 0]])
      >>> a[2, 2].execute()
      array([[0, 0, 0],
             [0, 0, 0],
             [0, 0, 4]])
  
      The wrap option affects only tall matrices:
  
      >>> # tall matrices no wrap
      >>> a = mt.zeros((5, 3), int)
      >>> mt.fill_diagonal(a, 4)
      >>> a.execute()
      array([[4, 0, 0],
             [0, 4, 0],
             [0, 0, 4],
             [0, 0, 0],
             [0, 0, 0]])
  
      >>> # tall matrices wrap
      >>> a = mt.zeros((5, 3), int)
      >>> mt.fill_diagonal(a, 4, wrap=True)
      >>> a.execute()
      array([[4, 0, 0],
             [0, 4, 0],
             [0, 0, 4],
             [0, 0, 0],
             [4, 0, 0]])
  
      >>> # wide matrices
      >>> a = mt.zeros((3, 5), int)
      >>> mt.fill_diagonal(a, 4, wrap=True)
      >>> a.execute()
      array([[4, 0, 0, 0, 0],
             [0, 4, 0, 0, 0],
             [0, 0, 4, 0, 0]])
  
      The anti-diagonal can be filled by reversing the order of elements
      using either `numpy.flipud` or `numpy.fliplr`.
  
      >>> a = mt.zeros((3, 3), int)
      >>> mt.fill_diagonal(mt.fliplr(a), [1,2,3])  # Horizontal flip
      >>> a.execute()
      array([[0, 0, 1],
             [0, 2, 0],
             [3, 0, 0]])
      >>> mt.fill_diagonal(mt.flipud(a), [1,2,3])  # Vertical flip
      >>> a.execute()
      array([[0, 0, 3],
             [0, 2, 0],
             [1, 0, 0]])
  
      Note that the order in which the diagonal is filled varies depending
      on the flip function.
      """
  
      if not isinstance(a, Tensor):
          raise TypeError(f"`a` should be a tensor, got {type(a)}")
      if a.ndim < 2:
          raise ValueError("array must be at least 2-d")
      if a.ndim > 2 and len(set(a.shape)) != 1:
          raise ValueError(
              "All dimensions of input must be of equal length"
          )
  
      # process val
      if isinstance(val, ENTITY_TYPE):
          val = astensor(val)
          if val.ndim > 1:
              val = val.ravel()
          val_input = val
      else:
          val = np.asarray(val)
          if val.ndim > 1:
              val = val.ravel()
          val_input = None
  
      op = TensorFillDiagonal(val=val, wrap=wrap, dtype=a.dtype)
      t = op(a, val=val_input)
      a.data = t.data
  
imports: import numpy
  from ...None import opcodes
  from ...core import ENTITY_TYPE, recursive_tile
  from ...serialization.serializables import KeyField, AnyField, BoolField, Int32Field
  from ...utils import has_unknown_shape, ceildiv
  from ..operands import TensorOperand, TensorOperandMixin
  from ..datasource import tensor
  from ..array_utils import as_same_device, device
  from ..core import Tensor, TENSOR_TYPE
  from ..utils import decide_unify_split
label_desc: Use image node_id 1 for calling a global function with example usage: fill_diagonal(a, val, wrap) without return types