int TensorInterface::NumDims() const { return tensor_.dims(); }

int64_t TensorInterface::Dim(int dim_index) const {
  return static_cast<int64_t>(tensor_.dim_size(dim_index));
}

int64_t TensorInterface::NumElements() const {
  return static_cast<int64_t>(tensor_.NumElements());
}

size_t TensorInterface::ByteSize() const {

  int64_t x_row_dim = x_shape[x_shape.size() - 2];
  int64_t x_col_dim = x_shape[x_shape.size() - 1];
  int64_t y_row_dim = y_shape[y_shape.size() - 2];
  int64_t y_col_dim = y_shape[y_shape.size() - 1];
  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

 public:
  ShapeInference() = default;
  explicit ShapeInference(ArrayRef<ArrayRef<int64_t>> input_shapes)
      : input_shapes_(input_shapes) {}
  void runOnOperation() override {
    // Parse `input_arg_shapes_` if provided (test only)
    SmallVector<ArrayRef<int64_t>> input_shapes_vec;
    absl::StatusOr<SmallVector<SmallVector<int64_t>>> parsed_shapes;
    if (!input_arg_shapes_.empty()) {

// fake quant op used in the training to use the nudged scale as well.
static void getNudgedScaleAndZeroPoint(int64_t qmin, int64_t qmax, double rmin,
                                       double rmax, double &scale,
                                       int64_t &nudgedZeroPoint) {
  // Determine the scale.
  const double qminDouble = qmin;
  const double qmaxDouble = qmax;

    // Thus, we fail to match if the consuming reshape rank is larger.
    ArrayRef<int64_t> input_shape = input_type.getShape();
    if (reshape_shape.size() > input_shape.size()) return failure();

    // Extend the input shape with leading 1s to match the broadcast shape.
    ArrayRef<int64_t> broadcast_shape = output_type.getShape();
    SmallVector<int64_t, 4> input_shape_extended;

}

namespace {

int64_t GetFirstIfIndicesAreContiguous(Value indices) {
  auto type = mlir::dyn_cast<RankedTensorType>(indices.getType());
  if (!type) return -1;
  auto indices_op = indices.getDefiningOp();
  if (!indices_op) return -1;
  auto const_op = llvm::dyn_cast<TF::ConstOp>(indices_op);
  if (!const_op) return -1;
  int64_t last_index = -1;
  int64_t first_index = -1;

inline bool HasRankAtMost(Value value, int64_t rank) {
  RankedTensorType type = GetRankedTensorTypeForOperand(value);
  return !type || type.getRank() <= rank;
}

inline bool IsUnknownDimOrRank(int64_t dim_or_rank) {
  return dim_or_rank == -1;
}

// Returns dimension index for the given TensorFlow axis that supports negative
// indexing.
inline int64_t GetDimForAxis(int64_t axis, int64_t rank) {

    MutableArrayRef<int64_t> input_shape, ArrayRef<int64_t> sparse_begin,
    ArrayRef<int64_t> sparse_end, ArrayRef<int64_t> sparse_strides,
    int32_t begin_mask, int32_t end_mask, int32_t ellipsis_mask,
    int32_t new_axis_mask, int32_t shrink_axis_mask,
    SmallVectorImpl<int64_t> *begin, SmallVectorImpl<int64_t> *end,
    SmallVectorImpl<int64_t> *stride) {
  int64_t num_sparse_indices = sparse_begin.size();

// pointer.
void CallFunctionsBase(::testing::benchmark::State& state) {
  int64_t sum = 0;
  typedef int64_t (*Func)(int64_t a, int64_t b);
  volatile Func f_raw = AddRaw;
  Func f = f_raw;
  size_t i = 0;
  for (auto dummy : state) {
    sum = f(sum, i);
    i++;
  }
  // volatile int64_t result = sum;
}

BENCHMARK(CallFunctions)->Arg(1 << 10);
BENCHMARK(CallFunctionsIndirect)->Arg(1 << 10);

    auto output = op->getResult(0);
    auto output_type =
        mlir::dyn_cast_or_null<mlir::RankedTensorType>(output.getType());
    if (!output_type || !output_type.hasStaticShape()) return false;

    *count = output_type.getNumElements();
    return true;
  }

  static bool GetInputTensorTotalSize(mlir::Operation* op, int64_t* count) {
    int64_t total_count = 0;