code: 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)
  
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 14 for calling a global function with example usage: _set_batch_non_diagonal(array, diag_val) without return types