<commit_before>import numpy

# TODO(okuta): Implement place


# TODO(okuta): Implement put


# TODO(okuta): Implement putmask


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, ..., i]`` all identical. This function
    modifies the input array in-place, it does not return a value.

    Args:
        a (cupy.ndarray): The array, at least 2-D.
        val (scalar): The value to be written on the diagonal.
            Its type must be compatible with that of the array a.
        wrap (bool): If specified, the diagonal is "wrapped" after N columns.
            This affects only tall matrices.

    Examples
    --------
    >>> a = cupy.zeros((3, 3), int)
    >>> cupy.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 = cupy.zeros((3, 3, 3, 3), int)
    >>> cupy.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 = cupy.zeros((5, 3),int)
    >>> cupy.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 = cupy.zeros((5, 3),int)
    >>> cupy.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 = cupy.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]])

    .. seealso:: :func:`numpy.fill_diagonal`
    """
    # The followings are imported from the original numpy
    if a.ndim < 2:
        raise ValueError("array must be at least 2-d")
    end = None
    if a.ndim == 2:
        step = a.shape[1] + 1
        if not wrap:
            end = a.shape[1] * a.shape[1]
    else:
        if not numpy.alltrue(numpy.diff(a.shape) == 0):
            raise ValueError("All dimensions of input must be of equal length")
        step = 1 + numpy.cumprod(a.shape[:-1]).sum()

    # we use a.ravel() instead of a.flat
    a.ravel()[:end:step] = val<commit_msg>Add reason why fill_diagonal() uses ravel()<commit_after>import numpy

# TODO(okuta): Implement place


# TODO(okuta): Implement put


# TODO(okuta): Implement putmask


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, ..., i]`` all identical. This function
    modifies the input array in-place, it does not return a value.

    Args:
        a (cupy.ndarray): The array, at least 2-D.
        val (scalar): The value to be written on the diagonal.
            Its type must be compatible with that of the array a.
        wrap (bool): If specified, the diagonal is "wrapped" after N columns.
            This affects only tall matrices.

    Examples
    --------
    >>> a = cupy.zeros((3, 3), int)
    >>> cupy.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 = cupy.zeros((3, 3, 3, 3), int)
    >>> cupy.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 = cupy.zeros((5, 3),int)
    >>> cupy.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 = cupy.zeros((5, 3),int)
    >>> cupy.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 = cupy.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]])

    .. seealso:: :func:`numpy.fill_diagonal`
    """
    # The followings are imported from the original numpy
    if a.ndim < 2:
        raise ValueError("array must be at least 2-d")
    end = None
    if a.ndim == 2:
        step = a.shape[1] + 1
        if not wrap:
            end = a.shape[1] * a.shape[1]
    else:
        if not numpy.alltrue(numpy.diff(a.shape) == 0):
            raise ValueError("All dimensions of input must be of equal length")
        step = 1 + numpy.cumprod(a.shape[:-1]).sum()

    # Since the current cupy does not support a.flat,
    # we use a.ravel() instead of a.flat
    a.ravel()[:end:step] = val
{'fraction_non_alphanumeric': 0.16989207294380348, 'fraction_numerical': 0.043542984741347227, 'mean_word_length': 1.868196371398079, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 125, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3913431', 'n_tokens_mistral': 2079, 'n_tokens_neox': 1903, 'n_words': 654}