code: def _non_diagonal_idx_array(batch_size, n):
      idx_offsets = np.arange(
          start=0, stop=batch_size * n * n, step=n * n, dtype=np.int32
      ).reshape((batch_size, 1))
      idx = (
          np.ravel_multi_index(np.tril_indices(n, -1), (n, n))
          .reshape((1, -1))
          .astype(np.int32)
      )
      return cuda.to_gpu(idx + idx_offsets)
  
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 4 for calling a global function with example usage: _non_diagonal_idx_array(batch_size, n) and returns: cuda