Function ResizeOutputTensor

tensorflow/lite/kernels/batch_to_space_nd.cc:56–99

const int kSpatialDimensionNum = 2;

TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
                                BatchToSpaceNDContext* op_context) {
  TfLiteIntArray* input_size = op_context->input->dims;
  const int* block_shape = GetTensorData<int32>(op_context->block_shape);
  const int* crops = GetTensorData<int32>(op_context->crops);

  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->block_shape),
                    kBlockSizeDimensionNum);
  TF_LITE_ENSURE_EQ(context, op_context->block_shape->dims->data[0],
                    kSpatialDimensionNum);
  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->crops),
                    kSpatialDimensionNum);

  TF_LITE_ENSURE(context, crops[0] >= 0);
  TF_LITE_ENSURE(context, crops[1] >= 0);
  TF_LITE_ENSURE(context, crops[2] >= 0);
  TF_LITE_ENSURE(context, crops[3] >= 0);

  // Number of batch must be multiple of (block_shape[0] * block_shape[1]).
  TF_LITE_ENSURE_EQ(context,
                    input_size->data[0] % (block_shape[0] * block_shape[1]), 0);

  const int output_batch_size =
      input_size->data[0] / (block_shape[0] * block_shape[1]);

  const int crops_top = crops[0];
  const int crops_bottom = crops[1];
  const int crops_left = crops[2];
  const int crops_right = crops[3];
  const int output_height =
      input_size->data[1] * block_shape[0] - crops_top - crops_bottom;
  const int output_width =
      input_size->data[2] * block_shape[1] - crops_left - crops_right;

  const int output_channel_size = input_size->data[3];

  TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_size);
  output_size->data[0] = output_batch_size;
  output_size->data[1] = output_height;
  output_size->data[2] = output_width;
  output_size->data[3] = output_channel_size;

  return context->ResizeTensor(context, op_context->output, output_size);
}

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);