code: def gather_nd(data_np, indices_np, batch_dims=0):
      """gather_nd implemented using numpy"""
      data_shape = data_np.shape
      indices_shape = indices_np.shape
  
      def gather_nd_batch_dims_1_ref(data, indices):
          res = []
          for i, row in enumerate(data):
              indices_tuple = tuple(
                  indices[:, i]
              )  # the indices for the i-th batch
              res.append(row[indices_tuple])
          # stack on the batch dim
          return np.stack(res, 0)
  
      if batch_dims > 1:
          data_np_reshape = np.reshape(
              data_np, (-1,) + data_shape[batch_dims:]
          )
          indices_np_reshape = np.reshape(
              indices_np,
              (indices_shape[0], -1)
              + indices_shape[(batch_dims + 1) :],
          )
  
          ref_res = gather_nd_batch_dims_1_ref(
              data_np_reshape, indices_np_reshape
          )
  
          out_shape = (
              indices_shape[1 : (batch_dims + 1)] + ref_res.shape[1:]
          )
          ref_res = np.reshape(ref_res, out_shape)
      elif batch_dims == 1:
          ref_res = gather_nd_batch_dims_1_ref(data_np, indices_np)
      else:
          ref_res = data_np[tuple(indices_np)]
  
      return ref_res
  
label_desc: Use image node_id 1 for calling a global function with example usage: gather_nd(data_np, indices_np, batch_dims) and returns: ref_res, np