100a265d7f0324bff1cce2f084cf23fc44bcc2c5,lower_triangular_matrix.py,,_set_batch_non_diagonal,#Any#Any#,29

Before Change

def _set_batch_non_diagonal(array, non_diag_val):
    batch_size, m, n = array.shape
    assert m == n
    rows, cols = np.tril_indices(n, -1)
    array[:, rows, cols] = non_diag_val

After Change

def _set_batch_non_diagonal(array, diag_val):
    xp = cuda.get_array_module(array)
    if xp == np:
        _set_batch_non_diagonal_cpu(array, diag_val)
    else:
        _set_batch_non_diagonal_gpu(array, diag_val)

Instances