text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from __future__ import unicode_literals
  
  import numpy as np  # type: ignore
  
  import onnx
  from ..base import Base
  from . import expect
  
  
  def gather_nd_impl(data, indices, batch_dims):
      # type: (np.ndarray, np.ndarray, int) -> np.ndarray
      # Note the data rank - will be reused multiple times later
      data_rank = len(data.shape)
  
      # Check input tensors' shape/rank condition
      assert indices.shape[-1] <= data_rank
  
      #The list of data/indice shape of batch_dims
      batch_dims_shape = []
  
      #The number of elements in the batch_dims for data/indice array
      batch_dims_size = 1
  
      # Check the shape of indice and data are identicial for batch dims.
      for i in range(batch_dims):
          batch_dims_shape.append(indices.shape[i])
          batch_dims_size *= indices.shape[i]
  
      # Compute output of the op as below
  
      # Compute shape of output array
      output_shape = batch_dims_shape + list(indices.shape)[batch_dims:-1] if (indices.shape[-1] == data_rank - batch_dims) \
       else batch_dims_shape + list(indices.shape)[batch_dims:-1] + list(data.shape)[batch_dims + indices.shape[-1]:]
  
      # Placeholder for output data
      output_data_buffer = []
  
      # Flatten 'indices' to 2D array
      reshaped_indices = indices.reshape(batch_dims_size, -1, indices.shape[-1])
  
      # Flatten 'data' to array of shape (batch_dim_size, data.shape[batch_dimes:])
      reshaped_data = data.reshape((batch_dims_size, ) + data.shape[batch_dims:])
  
      # gather each scalar value from 'data'
      for batch_dim in range(reshaped_indices.shape[0]):
          for outer_dim in range(reshaped_indices.shape[1]):
              gather_index = tuple(reshaped_indices[batch_dim][outer_dim])
              output_data_buffer.append(reshaped_data[(batch_dim,) + gather_index])
      return np.asarray(output_data_buffer, dtype=data.dtype).reshape(output_shape)
  
  
  class GatherND(Base):
  
      @staticmethod
      def export_int32():  # type: () -> None
          node = onnx.helper.make_node(
              'GatherND',
              inputs=['data', 'indices'],
              outputs=['output'],
          )
  
          data = np.array([[0, 1], [2, 3]], dtype=np.int32)
          indices = np.array([[0, 0], [1, 1]], dtype=np.int64)
          output = gather_nd_impl(data, indices, 0)
          expected_output = np.array([0, 3], dtype=np.int32)
          assert (np.array_equal(output, expected_output))
          expect(node, inputs=[data, indices], outputs=[output],
                 name='test_gathernd_example_int32')
  
      @staticmethod
      def export_float32():  # type: () -> None
          node = onnx.helper.make_node(
              'GatherND',
              inputs=['data', 'indices'],
              outputs=['output'],
          )
  
          data = np.array([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype=np.float32)
          indices = np.array([[[0, 1]], [[1, 0]]], dtype=np.int64)
          output = gather_nd_impl(data, indices, 0)
          expected_output = np.array([[[2, 3]], [[4, 5]]], dtype=np.float32)
          assert (np.array_equal(output, expected_output))
          expect(node, inputs=[data, indices], outputs=[output],
                 name='test_gathernd_example_float32')
  
      @staticmethod
      def export_int32_batchdim_1():  # type: () -> None
          node = onnx.helper.make_node(
              'GatherND',
              inputs=['data', 'indices'],
              outputs=['output'],
              batch_dims=1,
          )
  
          data = np.array([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], dtype=np.int32)
          indices = np.array([[1], [0]], dtype=np.int64)
          output = gather_nd_impl(data, indices, 1)
          expected_output = np.array([[2, 3], [4, 5]], dtype=np.int32)
          assert (np.array_equal(output, expected_output))
          expect(node, inputs=[data, indices], outputs=[output],
                 name='test_gathernd_example_int32_batch_dim1')
  