def fill_diagonal(a, val, wrap=False):
if a.ndim < <NUM_LIT:2>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>end = None<EOL>if a.ndim == <NUM_LIT:2>:<EOL><INDENT>step = a.shape[<NUM_LIT:1>] + <NUM_LIT:1><EOL>if not wrap:<EOL><INDENT>end = a.shape[<NUM_LIT:1>] * a.shape[<NUM_LIT:1>]<EOL><DEDENT><DEDENT>else:<EOL><INDENT>if not alltrue(diff(a.shape) == <NUM_LIT:0>):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>step = <NUM_LIT:1> + (cumprod(a.shape[:-<NUM_LIT:1>])).sum()<EOL><DEDENT>a.flat[:end:step] = val<EOL>
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, ..., 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 affect 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]])

    # tall matrices no wrap
    >>> a = np.zeros((5, 3),int)
    >>> fill_diagonal(a, 4)
    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)
    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)
    array([[4, 0, 0, 0, 0],
           [0, 4, 0, 0, 0],
           [0, 0, 4, 0, 0]])