def nzcol(mat, row_idx):
    """Return the nonzero elements of given row in a CSR matrix.

    Parameters
    ----------
    mat : CSR matrix
        Input matrix.
    row_idx : int
        The index of the row (if `mat` is CSR) for which the nonzero
        elements are desired.

    Returns
    -------
    nz : array of int
        The location of nonzero elements of `mat[main_axis_idx]`.

    Examples
    --------
    >>> mat = sparse.csr_matrix(np.array([[0, 1, 0, 0], [0, 5, 8, 0]]))
    >>> nzcol(mat, 1)
    array([1, 2], dtype=int32)
    >>> mat[1, 2] = 0
    >>> nzcol(mat, 1)
    array([1], dtype=int32)
    """
    return mat[row_idx].nonzero()[1]