code: def fill_diagonal(a, val, wrap=False):
      """Fill the main diagonal of the given array of any dimensionality.
  
      For an array `a` with ``a.ndim >= 2``, the diagonal is the list of
      locations with indices ``a[i, ..., i]`` all identical. This function
      modifies the input array in-place, it does not return a value.
  
      Parameters
      ----------
      a : array, at least 2-D.
        Array 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 array 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
      -----
      .. versionadded:: 1.4.0
  
      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
      --------
      >>> a = np.zeros((3, 3), int)
      >>> np.fill_diagonal(a, 5)
      >>> a
      array([[5, 0, 0],
             [0, 5, 0],
             [0, 0, 5]])
  
      The same function can operate on a 4-D array:
  
      >>> a = np.zeros((3, 3, 3, 3), int)
      >>> np.fill_diagonal(a, 4)
  
      We only show a few blocks for clarity:
  
      >>> a[0, 0]
      array([[4, 0, 0],
             [0, 0, 0],
             [0, 0, 0]])
      >>> a[1, 1]
      array([[0, 0, 0],
             [0, 4, 0],
             [0, 0, 0]])
      >>> a[2, 2]
      array([[0, 0, 0],
             [0, 0, 0],
             [0, 0, 4]])
  
      The wrap option affects only tall matrices:
  
      >>> # tall matrices no wrap
      >>> a = np.zeros((5, 3),int)
      >>> fill_diagonal(a, 4)
      >>> a
      array([[4, 0, 0],
             [0, 4, 0],
             [0, 0, 4],
             [0, 0, 0],
             [0, 0, 0]])
  
      >>> # tall matrices wrap
      >>> a = np.zeros((5, 3),int)
      >>> fill_diagonal(a, 4, wrap=True)
      >>> a
      array([[4, 0, 0],
             [0, 4, 0],
             [0, 0, 4],
             [0, 0, 0],
             [4, 0, 0]])
  
      >>> # wide matrices
      >>> a = np.zeros((3, 5),int)
      >>> fill_diagonal(a, 4, wrap=True)
      >>> a
      array([[4, 0, 0, 0, 0],
             [0, 4, 0, 0, 0],
             [0, 0, 4, 0, 0]])
  
      """
      if a.ndim < 2:
          raise ValueError("array must be at least 2-d")
      end = None
      if a.ndim == 2:
          # Explicit, fast formula for the common case.  For 2-d arrays, we
          # accept rectangular ones.
          step = a.shape[1] + 1
          # This is needed to don't have tall matrix have the diagonal wrap.
          if not wrap:
              end = a.shape[1] * a.shape[1]
      else:
          # For more than d=2, the strided formula is only valid for arrays with
          # all dimensions equal, so we check first.
          if not alltrue(diff(a.shape) == 0):
              raise ValueError(
                  "All dimensions of input must be of equal length"
              )
          step = 1 + (cumprod(a.shape[:-1])).sum()
  
      # Write the value out into the diagonal.
      a.flat[:end:step] = val
  
imports: from __future__ import division, absolute_import, print_function
  import functools
  import sys
  import math
  import numpy.core.numeric
  from numpy.core.numeric import asarray, ScalarType, array, alltrue, cumprod, arange, ndim
  from numpy.core.numerictypes import find_common_type, issubdtype
  import numpy.matrixlib
  from .function_base import diff
  from numpy.core.multiarray import ravel_multi_index, unravel_index
  from numpy.core.overrides import set_module
  from numpy.core import overrides, linspace
  from numpy.lib.stride_tricks import as_strided
label_desc: Use image node_id 4 for calling a global function with example usage: fill_diagonal(a, val, wrap) without return types