std::swap(rhs_shape[rhs_dims - 1], rhs_shape[rhs_dims - 2]);
  }

  const int64_t rows = lhs_shape[lhs_dims - 2];
  const int64_t cols = rhs_shape[rhs_dims - 1];

  if (lhs_shape[lhs_dims - 1] != rhs_shape[rhs_dims - 2]) {
    // Input dimensions must be compatible for multiplication.
    return failure();
  }

  const auto matmul_type = RankedTensorType::get({rows, cols}, element_type);

  bool isSupportedAxis(mlir::Value value, int64_t rank) const {
    auto const_op =
        dyn_cast_or_null<mlir::arith::ConstantOp>(value.getDefiningOp());
    if (!const_op) {
      return false;
    }
    auto axes = dyn_cast<DenseIntElementsAttr>(const_op.getValueAttr());
    if (!axes || axes.getNumElements() != 1) {
      return false;
    }

  // CHECK-NEXT: [[PREFIX:%.*]] = tf_mlrt.constop
  %cst = "tf.Const"() {__op_key = 0: i32, value = dense<"restore_ariables"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
  // CHECK-NEXT: [[SLICE:%.*]] = tf_mlrt.constop
  %cst_0 = "tf.Const"()  {__op_key = 1: i32, value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
  // CHECK-NEXT: [[NAME:%.*]] = tf_mlrt.constop

 public:
  explicit CastBf16OpsToF32(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}

 private:
  LogicalResult match(Operation* op) const override {
    if (isa<TF::CastOp, TF::ConstOp>(op) ||
        op->getName().hasTrait<OpTrait::ZeroOperands>()) {
      return failure();
    }
    for (Value input : op->getOperands()) {
      if (getElementTypeOrSelf(input).isBF16()) {

  explicit QuantizeCompositeFunctionsPass() = default;

  explicit QuantizeCompositeFunctionsPass(
      const QuantMethod quantization_method, const OpSet target_opset,
      const bool enable_per_channel_quantization,
      const int min_num_elements_for_weights,
      const bool enable_legacy_weight_only,
      std::optional<const std::string> mlir_dump_file_name)
      : enable_legacy_weight_only_(enable_legacy_weight_only),

            bvector.set(use.getOperandNumber());
        }
      }

      // Replace the arg with a tf.Const op in the function body.
      builder.setInsertionPointToStart(&func.getBody().front());
      auto const_op = builder.create<TF::ConstOp>(global_tensor.getLoc(),
                                                  *global_tensor.getValue());
      args_to_erase.set(val.getArgNumber());

    auto inserted_value = values->insert(operand).second;
    if (inserted_value) continue;
    // We can only clone the constant op or const->dequantize combo. The latter
    // case is useful for float16 quantization. Since all ops have been
    // legalized to tflite ops, so we only care about ConstOp or QConstOp or
    // mlir constant op.
    Operation* input_op = operand.getDefiningOp();
    if (input_op == nullptr) return failure();

    llvm::SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
  const int64_t num_spatial_dims = std::is_same<OpT, Conv2DOpAdaptor>() ? 2 : 3;
  const int64_t num_dims = 2 + num_spatial_dims;
  const Value input = op.getInput();
  const Value filter = op.getFilter();
  const TensorType input_ty = mlir::cast<TensorType>(input.getType());
  const TensorType filter_ty = mlir::cast<TensorType>(filter.getType());

    IntegerAttr group_key = rewriter.getI32IntegerAttr(0);

    auto const_group_size = rewriter.create<TF::ConstOp>(
        assign_group->getLoc(), assign_group.getResult(0).getType(),
        group_size);
    auto const_group_key = rewriter.create<TF::ConstOp>(
        assign_group->getLoc(), assign_group.getResult(1).getType(), group_key);

    return rewriter.createOrFold<TF::CastOp>(loc, new_type, val,
                                             rewriter.getBoolAttr(false));
  }
  return rewriter.createOrFold<TF::CastOp>(
      loc, UnrankedTensorType::get(new_ele_type), val,
      rewriter.getBoolAttr(false));
}

// Utility function to-
// 1. Create a tfl.const op with an int32_t values, from an MLIR Value, if the