code: def _get_batch_diagonal_cpu(array):
      batch_size, m, n = array.shape
      assert m == n
      rows, cols = np.diag_indices(n)
      return array[:, rows, cols]
  
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 1 for calling a global function with example usage: _get_batch_diagonal_cpu(array) and returns: array