#include "tensorflow/core/platform/test.h"

namespace mlir {
namespace TFL {
namespace {

TEST(SizeUtilTest, TestConvertsSize) {
  ASSERT_EQ(ConvertToTfliteSize(1), 1);
  ASSERT_EQ(ConvertToTfliteSize(-1), -1);
  ASSERT_EQ(ConvertToTfliteSize(mlir::ShapedType::kDynamic), -1);
}

}  // namespace
}  // namespace TFL

#include "tensorflow/compiler/mlir/lite/utils/size_utils.h"

#include <cstdint>

#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project

namespace mlir {
namespace TFL {

int32_t ConvertToTfliteSize(int64_t size) {
  return mlir::ShapedType::isDynamic(size) ? -1 : static_cast<int32_t>(size);
}

}  // namespace TFL

namespace mlir {
namespace TFL {

// Converts a TF size (64-bit) to TFLite (32-bit) and properly converts TF's
// value for dynamic size (`std::numeric_limits<int64_t>::min()`) to the
// TFLite-specific value.
int32_t ConvertToTfliteSize(int64_t size);

}  // namespace TFL
}  // namespace mlir

    SmallVector<Attribute, 4> result_shape_data(4);
    for (int i = 0; i < 4; ++i) {
      auto size = result_shape[i];
      result_shape_data[i] =
          rewriter.getI32IntegerAttr(ConvertToTfliteSize(size));
    }
    auto shape_attr = DenseElementsAttr::get(shape_type, result_shape_data);
    auto shape = rewriter.create<TF::ConstOp>(loc, shape_type, shape_attr);

      mlir::SmallVector<mlir::Attribute, 4> shape;
      for (auto s : new_shape) {
        shape.push_back(
            builder.getI32IntegerAttr(mlir::TFL::ConvertToTfliteSize(s)));
      }
      auto output_shape = DenseElementsAttr::get(shape_type, shape);
      auto shape_op = builder.create<tfl::ConstOp>(loc, output_shape);
      op_state.addOperands({shape_op});
    }

    // This is to workaround the unnecessary cast i64 -> i32.
    SmallVector<int32_t, 4> new_shape_array;
    for (auto size : output_type.getShape()) {
      new_shape_array.push_back(ConvertToTfliteSize(size));
    }

    auto new_shape = rewriter.create<TFL::ConstOp>(
        loc, DenseIntElementsAttr::get(
                 tensorflow::GetTypeFromTFTensorShape(