// element is arbitrarily considered to be 4 bytes.
constexpr int64_t kAssumedNumBytesPerElem = 4;

int64_t GetSizeOfIntOrFloatConst(TF::ConstOp const_op) {
  const Type dtype = const_op.getDtype();
  const ElementsAttr const_value = const_op.getValue();

  const auto bytes_per_elem =
      static_cast<int64_t>(dtype.getIntOrFloatBitWidth() / CHAR_BIT);

  return bytes_per_elem * const_value.getNumElements();

        scatter_op, "update_window_dims are not leading or trailing indices");

  SmallVector<int64_t, 4> permutation_array(updates_type.getRank());
  int64_t dim = 0;
  // Move leading indices to the back of the array.
  const auto permutation_array_size = permutation_array.size();
  for (int64_t i = update_window_dims.size(); i < permutation_array_size; ++i) {
    permutation_array[i] = dim;
    ++dim;
  }

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 {

// with `size` number of values.
bool IsIotaAttr(ArrayRef<int64_t> arr, int64_t size);

// Returns a DenseIntElementsAttr for a permutation and the shape after
// applying the permutation to a given shape through a transpose.
PermutationAndShape GetPermutationAndTransposedShape(
    llvm::ArrayRef<int64_t> permutation_array, ShapedType input_type,
    ConversionPatternRewriter& rewriter);

 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()) {

}

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) {

  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];

// 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;