absl::Status UnfreezeConstantsAndSaveVariables(
    const absl::string_view checkpoint_dir, mlir::MLIRContext &ctx,
    mlir::ModuleOp module_op) {
  TF_RETURN_IF_ERROR(RunPasses(
      /*name=*/kTfQuantConstantUnfreezingStepName, /*add_passes_func=*/
      [](mlir::PassManager &pm) {
        pm.addPass(mlir::quant::CreateUnfreezeConstantsPass());
      },
      ctx, module_op));

  if (const absl::Status create_dir_status =

      if (i != feature_index) {
        reduce_dim_axes.push_back(i);
      }
    }
    auto mean_dims = DenseIntElementsAttr::get(dims_type, reduce_dim_axes);
    ::mlir::TF::ConstOp reduce_dim_op =
        rewriter.create<TF::ConstOp>(bn_op.getLoc(), mean_dims);
    int64_t feature_dim_size = input_type.getDimSize(feature_index);
    auto mean_var_type = RankedTensorType::get(/*shape=*/{feature_dim_size},

      state.padding = rewriter.getStringAttr("VALID");
    }
    auto conv_op = static_cast<const ConcreteType *>(this)->createTFLOp(
        &state, rewriter, op->getLoc(), tf_op.getType(), input, filter, bias);

    rewriter.replaceOp(op, conv_op.getResult());
    return success();
  }

  const IntegerAttr intAttrOne;

 private:
  bool allow_bf16_and_f16_type_legalization_;
};

}

TF::ConstOp Create1DConstantOp(const std::vector<int>& value, Location loc,
                               OpBuilder* builder) {
  auto type =
      mlir::RankedTensorType::get(value.size(), builder->getIntegerType(32));
  auto dense_values = mlir::DenseIntElementsAttr::get(type, value);
  return builder->create<TF::ConstOp>(loc, dense_values);
}

TF::ConstOp CreateScalarConstantOp(int value, Location loc,

  }

  const auto type = RankedTensorType::get(shape, builder.getF32Type());
  const auto value_attr = DenseFPElementsAttr::get(type, values);
  return builder.create<TF::ConstOp>(loc, value_attr);
}

// Creates a 1D array with integer/float type.
template <typename T>
Value Create1DConstValue(OpBuilder& builder, const Location loc,
                         const SmallVector<T>& values) {

      auto qcast = b.create<quantfork::QuantizeCastOp>(
          q.getLoc(), q.getOutput().getType(), q.getInput());
      q.getOutput().replaceAllUsesWith(qcast);
      q.erase();
    } else if (auto q = llvm::dyn_cast<ConstOp>(op)) {
      auto value = q.getValue();
      auto type = q.getResult().getType();
      if (arith::ConstantOp::isBuildableWith(value, type)) {
        auto c = b.create<arith::ConstantOp>(q.getLoc(), q.getValue());

    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

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