TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  MicroContext* micro_context = GetMicroContext(context);

  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  TfLiteTensor* params = micro_context->AllocateTempInputTensor(node, kParams);
  TF_LITE_ENSURE(context, params != nullptr);
  TfLiteTensor* indices =
      micro_context->AllocateTempInputTensor(node, kIndices);
  TF_LITE_ENSURE(context, indices != nullptr);
  TfLiteTensor* output =
      micro_context->AllocateTempOutputTensor(node, kOutputTensor);
  TF_LITE_ENSURE(context, output != nullptr);

  switch (params->type) {
    case kTfLiteFloat32:
    case kTfLiteInt8:
      break;
    default:
      TF_LITE_KERNEL_LOG(context,
                         "Params of type '%s' are not supported by gather_nd.",
                         TfLiteTypeGetName(params->type));
      return kTfLiteError;
      break;
  }
  switch (indices->type) {
    case kTfLiteInt32:
      break;
    default:
      TF_LITE_KERNEL_LOG(context,
                         "Indices of type '%s' are not supported by gather_nd.",
                         TfLiteTypeGetName(indices->type));
      return kTfLiteError;
  }

  const int params_rank = NumDimensions(params);
  const int indices_rank = NumDimensions(indices);
  const int indices_nd = SizeOfDimension(indices, indices_rank - 1);
  if (params_rank < 1) {
    TF_LITE_KERNEL_LOG(context, "Params must be at least a vector.");
    return kTfLiteError;
  }
  if (indices_rank < 1) {
    TF_LITE_KERNEL_LOG(context, "Indices must be at least a vector.");
    return kTfLiteError;
  }
  if (indices_nd > params_rank) {
    TF_LITE_KERNEL_LOG(
        context, "Index innermost dimension length must be <= params rank.");
    return kTfLiteError;
  }
  if (indices_nd > MAX_INDICES_ND) {
    TF_LITE_KERNEL_LOG(context,
                       "Index innermost dimension length must not exceed %d.",
                       MAX_INDICES_ND);
    return kTfLiteError;
  }

  // Assign to output the input type.
  output->type = params->type;

  // TFLM gather_nd does not create the output tensor, but it needs to ensure
  // that the output shape is correct. The result shape is
  // indices.shape[:-1] + params.shape[indices.shape[-1]:]
  TfLiteIntArray* output_shape = output->dims;
  int output_index = 0;
  for (int i = 0; i < indices_rank - 1; ++i) {
    output_shape->data[output_index++] = indices->dims->data[i];
  }
  for (int i = indices_nd; i < params_rank; ++i) {
    output_shape->data[output_index++] = params->dims->data[i];
  }
  output_shape->size = output_index;

  micro_context->DeallocateTempTfLiteTensor(params);
  micro_context->DeallocateTempTfLiteTensor(indices);
  micro_context->DeallocateTempTfLiteTensor(output);
  return kTfLiteOk;
}
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  MicroContext* micro_context = GetMicroContext(context);

  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  TfLiteTensor* params = micro_context->AllocateTempInputTensor(node, kParams);
  TF_LITE_ENSURE(context, params != nullptr);
  TfLiteTensor* indices =
      micro_context->AllocateTempInputTensor(node, kIndices);
  TF_LITE_ENSURE(context, indices != nullptr);
  TfLiteTensor* output =
      micro_context->AllocateTempOutputTensor(node, kOutputTensor);
  TF_LITE_ENSURE(context, output != nullptr);

  switch (params->type) {
    case kTfLiteFloat32:
    case kTfLiteInt8:
      break;
    default:
      TF_LITE_KERNEL_LOG(context,
                         "Params of type '%s' are not supported by gather_nd.",
                         TfLiteTypeGetName(params->type));
      return kTfLiteError;
      break;
  }
  switch (indices->type) {
    case kTfLiteInt32:
      break;
    default:
      TF_LITE_KERNEL_LOG(context,
                         "Indices of type '%s' are not supported by gather_nd.",
                         TfLiteTypeGetName(indices->type));
      return kTfLiteError;
  }

  const int params_rank = NumDimensions(params);
  const int indices_rank = NumDimensions(indices);
  const int indices_nd = SizeOfDimension(indices, indices_rank - 1);
  if (params_rank < 1) {
    TF_LITE_KERNEL_LOG(context, "Params must be at least a vector.");
    return kTfLiteError;
  }
  if (indices_rank < 1) {
    TF_LITE_KERNEL_LOG(context, "Indices must be at least a vector.");
    return kTfLiteError;
  }
  if (indices_nd > params_rank) {
    TF_LITE_KERNEL_LOG(
        context, "Index innermost dimension length must be <= params rank.");
    return kTfLiteError;
  }
  if (indices_nd > MAX_INDICES_ND) {
    TF_LITE_KERNEL_LOG(context,
                       "Index innermost dimension length must not exceed %d.",
                       MAX_INDICES_ND);
    return kTfLiteError;
  }

   
  output->type = params->type;

   
   
   
  TfLiteIntArray* output_shape = output->dims;
  int output_index = 0;
  for (int i = 0; i < indices_rank - 1; ++i) {
    output_shape->data[output_index++] = indices->dims->data[i];
  }
  for (int i = indices_nd; i < params_rank; ++i) {
    output_shape->data[output_index++] = params->dims->data[i];
  }
  output_shape->size = output_index;

  micro_context->DeallocateTempTfLiteTensor(params);
  micro_context->DeallocateTempTfLiteTensor(indices);
  micro_context->DeallocateTempTfLiteTensor(output);
  return kTfLiteOk;
}