auto output_scale_type =
        mlir::dyn_cast<ShapedType>(op->getOperand(3).getType());
    if (!output_scale_type) {
      return failure();
    }
    if (output_scale_type.hasRank() && 0 < output_scale_type.getRank()) {
      output_quantization_axis = activation_quantization_axis;
    }
  }
  // For per-axis -> per-axis requantization, input and output quantization
  // axis must be equal.

}

// Returns true if the rank of the value equals to the given rank.
bool RankEquals(Value value, int rank) {
  auto rank_type = mlir::dyn_cast<RankedTensorType>(value.getType());
  return (rank_type && rank_type.getRank() == rank);
}

#include "tensorflow/compiler/mlir/lite/quantization/tensorflow/fallback_to_flex_patterns.inc"

void FallbackToFlexOps::runOnOperation() {
  if (mode_.empty()) return;

    // the first `TensorListSetItemOp`.
    if (auto shaped_type = element_shape.getType().dyn_cast<ShapedType>()) {
      if (shaped_type.hasRank() && shaped_type.getRank() == 0) {
        bool element_shape_acquired = false;
        auto uses = op.getResult().getUses();
        for (auto &use : llvm::make_early_inc_range(uses)) {
          if (TF::TensorListSetItemOp set_op =

  auto per_slice_shape = llvm::to_vector<8>(buffer_type.getShape());
  per_slice_shape[0] = 1;
  auto slice_sizes = GetR1Const(per_slice_shape, builder, loc);
  llvm::SmallVector<int64_t, 8> starts_in_update(buffer_type.getRank(), 0);
  for (int64_t i = 0; i < updates_type.getDimSize(0); ++i) {
    auto index = builder.create<TF::SliceOp>(
        loc, ArrayRef<Type>{GetSizeType(builder)},

// neither.
mlir::RankedTensorType GetBufferType(mlir::Type ty) {
  auto ranked_ty = mlir::dyn_cast_or_null<mlir::RankedTensorType>(ty);
  if (!ranked_ty) return {};

  int64_t rank = ranked_ty.getRank();
  llvm::SmallVector<int64_t, 4> dims = llvm::to_vector<4>(ranked_ty.getShape());
  auto encoding = mlir::dyn_cast_or_null<mlir::mhlo::TypeExtensionsAttr>(
      ranked_ty.getEncoding());

  if (ranked_ty.hasStaticShape()) return true;

  auto encoding =
      mlir::dyn_cast_or_null<TypeExtensionsAttr>(ranked_ty.getEncoding());
  if (!encoding) return false;

  for (int i = 0; i < ranked_ty.getRank(); ++i) {
    if (ranked_ty.isDynamicDim(i) &&
        encoding.getBounds()[i] == ShapedType::kDynamic) {
      return false;
    }
  }
  return true;
}

bool HasSymbolRefAttr(Operation* op) {

  }

  // Allow dynamic width and height dimensions only.
  auto result_ty = mlir::cast<TensorType>(op.getResult().getType());
  if (!result_ty.hasRank() || result_ty.getRank() != 4 ||
      result_ty.isDynamicDim(0) || result_ty.isDynamicDim(3)) {
    return rewriter.notifyMatchFailure(
        op, "only dynamic width and height dimensions are allowed");
  }

    // verify shardings that actually break a tensor apart.
    return mlir::success();
  }
  if (RankedTensorType ranked_type = mlir::dyn_cast<RankedTensorType>(type)) {
    const int64_t tensor_rank = ranked_type.getRank();
    int tile_assignment_rank = sharding->tile_assignment_dimensions_size();

    // When a tensor is partial or subgroup tiled, its tile assignment will

    if (bias_op != nullptr) {
      Type bias_type = bias_op->getResult(0).getType();
      if (bias_type != builder_.getNoneType()) {
        const int bias_rank = mlir::dyn_cast<ShapedType>(bias_type).getRank();
        adjusted_quant_dim = bias_rank > 1 ? bias_rank - 1 : 0;
      }
    }
  }

  for (const int non_bias_operand_index : non_bias_operand_indices) {
    const QuantState& non_bias_state =

  if (*count <= 0) return split.emitOpError("non-positive split count");
  auto buffer_type = split.getValue().getType().dyn_cast<RankedTensorType>();
  if (!buffer_type || !buffer_type.hasStaticShape() ||
      buffer_type.getRank() < 1) {
    return split.emitOpError("unknown or invalid split tensor shape");
  }
  int64_t length = buffer_type.getDimSize(0) / *count;
  for (const auto& len : lengths_const.getValue().getValues<APInt>()) {