docstr: Checks that batch_dims matches multiple calls to tf.gather().
code: def testBatchDimsMatchesPythonBatching(
      self, params_shape, indices_shape, batch_dims, axis, output_shape
  ):
      """Checks that batch_dims matches multiple calls to tf.gather()."""
      # Generate a `params` tensor with the indicated shape.
      params_size = np.prod(params_shape)
      params = np.reshape(np.arange(params_size), params_shape)
  
      # Generate an `indices` tensor with the indicated shape, where each index
      # is within the appropriate range.
      indices_size = np.prod(indices_shape)
      indices = np.reshape(np.arange(indices_size), indices_shape)
      indices = indices % params_shape[axis]
  
      # Perform repeated (batched) gather operations with numpy, to find the
      # expected result.
      expected = self._batchNumpyGather(
          params, indices, axis, batch_dims
      )
  
      # On Windows, we get an exception if we pass in the transformed numpy
      # arrays ("Failed to convert numpy ndarray to a Tensor (Unsupported
      # feed type)."); so convert them back to lists before calling tf.gather.
      params = params.tolist()
      indices = indices.tolist()
  
      result = array_ops.gather(
          params, indices, axis=axis, batch_dims=batch_dims
      )
      self.assertAllEqual(output_shape, result.shape.as_list())
      self.assertAllEqual(expected, result)
  
      # Run the same test for strings.
      params = _to_str_elements(params)
      expected = _to_str_elements(expected.tolist())
      result = array_ops.gather(
          params, indices, axis=axis, batch_dims=batch_dims
      )
  
      self.assertAllEqual(output_shape, result.shape.as_list())
      self.assertAllEqual(expected, result)
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import indexed_slices
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gradient_checker_v2
  from tensorflow.python.ops import gradients_impl
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import ref_variable
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 17 for calling the GatherTest obj's underlying member method code with example usage: obj.testBatchDimsMatchesPythonBatching(params_shape, indices_shape, batch_dims, axis, output_shape) without return types