code: def _set_batch_diagonal_gpu(array, diag_val):
      batch_size, m, n = array.shape
      assert m == n
      cuda.cupy.ElementwiseKernel(
          "T val, int32 idx",
          "raw T mat",
          "mat[idx] = val",
          "lower_triangular_matrix_set_diag",
      )(diag_val, _diagonal_idx_array(batch_size, n), array)
  
imports: from functools import lru_cache
  from chainer import cuda
  from chainer import function
  from chainer.utils import type_check
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: _set_batch_diagonal_gpu(array, diag_val) without return types