code: def test_diag_part(self):
      for x in (
          random_ops.random_uniform([3, 2, 2]),
          random_ops.random_uniform([3, 4, 2, 4, 2]),
      ):
          # pylint: disable=cell-var-from-loop
          def loop_fn(i):
              inp = array_ops.gather(
                  x, i
              )  # pylint: disable=redefined-builtin
              return array_ops.diag_part(inp)
  
          # pylint: disable=cell-var-from-loop
          self._test_loop_fn(loop_fn, 3)
  
imports: from tensorflow.python.eager import backprop
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import tensor_array_grad
  from tensorflow.python.ops.parallel_for import control_flow_ops
  from tensorflow.python.ops.parallel_for.test_util import PForTestCase
  from tensorflow.python.platform import test
label_desc: Use image node_id 37 for calling the ArrayTest obj's underlying member method code with example usage: obj.test_diag_part() and returns: array_ops