class_docstr: The iOS16 version.
  This section documents only the differences between this version and the
  iOS 15 :py:class:`~.iOS15.scatter_gather.gather_nd`.
  
  This version supports ``batch_dims``.
  Input parameter ``indices`` now supports ``int16`` and ``uint16``.
  
  Parameters
  ----------
  x: tensor<\*D, T> (Required)
  indices: tensor<\*K, I> (Required)
  batch_dims: const i32 (Optional. Default=``0``)
      * The number of batch dimensions.
  
  Returns
  -------
  tensor<\*V, T>
      * ``V = K[:-1] + D[batch_dims + K[-1]:]``, where ``D = x.shape`` and ``K = indices.shape``.
  
  Attributes
  ----------
  T: fp16, fp32, i32
  I: uint16, int16, int32
  
  References
  ----------
  See `tf.gather_nd <https://www.tensorflow.org/api_docs/python/tf/gather_nd>`_.
class_docstr_tok: ["The","iOS16","version",".","This","section","documents","only","the","differences","between","this","version","and","the","iOS","15",":","py",":","class",":","`","~.iOS15.scatter_gather.gather_nd","`",".","This","version","supports","``","batch_dims","``",".","Input","parameter","``","indices","``","now","supports","``","int16","``","and","``","uint16","``",".","Parameters","--","--","--","--","--","x",":","tensor","<","\\","*","D",",","T",">","(","Required",")","indices",":","tensor","<","\\","*","K",",","I",">","(","Required",")","batch_dims",":","const","i32","(","Optional",".","Default=","``","0","``",")","*","The","number","of","batch","dimensions",".","Returns","--","--","--","-","tensor","<","\\","*","V",",","T",">","*","``","V","=","K","[",":","-1","]","+","D","[","batch_dims","+","K","[","-1","]",":","]","``",",","where","``","D","=","x.shape","``","and","``","K","=","indices.shape","``",".","Attributes","--","--","--","--","--","T",":","fp16",",","fp32",",","i32","I",":","uint16",",","int16",",","int32","References","--","--","--","--","--","See","`","tf.gather_nd","<","https",":","\/\/www.tensorflow.org\/api_docs\/python\/tf\/gather_nd",">","`","_","."]
code: def type_inference(self):
      batch_dims = self.batch_dims.val
      indices_depth = self.indices.shape[-1]
      if indices_depth > self.x.rank - batch_dims:
          msg = "For node {}, indices.shape[-1] ({}) + batch_dims ({}) must be smaller or equal to the input rank {}".format(
              self.name, indices_depth, batch_dims, self.x.rank
          )
          raise ValueError(msg)
      out_type = self.x.dtype
      out_shape = (
          self.indices.shape[:-1]
          + self.x.shape[batch_dims + indices_depth :]
      )
      return types.tensor(out_type, out_shape)
  
file_path: coremltools/coremltools/converters/mil/mil/ops/defs/iOS16/scatter_gather.py
imports: from coremltools.converters.mil.mil import Operation, types
  from coremltools.converters.mil.mil.input_type import DefaultInputs, InputSpec, TensorInputType
  from coremltools.converters.mil.mil.operation import SYMBOL, VALUE, precondition
  from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op
  from coremltools.converters.mil.mil.ops.defs._utils import compute_gather
  from coremltools.converters.mil.mil.ops.defs.iOS15.scatter_gather import gather_along_axis
  from coremltools.converters.mil.mil.ops.defs.iOS16 import _IOS16_TARGET
label_desc: Use image node_id 2 for calling the gather_nd obj's underlying member method code with example usage: obj.type_inference() and returns: types