void FillDiagHelper(const TfLiteTensor* input, const TfLiteTensor* diag,
                    TfLiteTensor* output) {
  const int num_output_dims = output->dims->size;
  int batch_size = 1;
  for (int i = 0; i < num_output_dims - 2; ++i) {
    batch_size *= output->dims->data[i];
  }

  const int row_size = output->dims->data[num_output_dims - 2];
  const int col_size = output->dims->data[num_output_dims - 1];
  switch (output->type) {
    case kTfLiteInt64: {
      return FillDiag<int64_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    case kTfLiteInt32: {
      return FillDiag<int32_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    case kTfLiteInt16: {
      return FillDiag<int16_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    case kTfLiteInt8: {
      return FillDiag<int8_t>(input, diag, output, batch_size, row_size,
                              col_size);
    }
    case kTfLiteUInt8: {
      return FillDiag<uint8_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    default:
      return FillDiag<float>(input, diag, output, batch_size, row_size,
                             col_size);
  }
}
void FillDiagHelper(const TfLiteTensor* input, const TfLiteTensor* diag,
                    TfLiteTensor* output) {
  const int num_output_dims = output->dims->size;
  int batch_size = 1;
  for (int i = 0; i < num_output_dims - 2; ++i) {
    batch_size *= output->dims->data[i];
  }

  const int row_size = output->dims->data[num_output_dims - 2];
  const int col_size = output->dims->data[num_output_dims - 1];
  switch (output->type) {
    case kTfLiteInt64: {
      return FillDiag<int64_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    case kTfLiteInt32: {
      return FillDiag<int32_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    case kTfLiteInt16: {
      return FillDiag<int16_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    case kTfLiteInt8: {
      return FillDiag<int8_t>(input, diag, output, batch_size, row_size,
                              col_size);
    }
    case kTfLiteUInt8: {
      return FillDiag<uint8_t>(input, diag, output, batch_size, row_size,
                               col_size);
    }
    default:
      return FillDiag<float>(input, diag, output, batch_size, row_size,
                             col_size);
  }
}