}

  auto scalar_type = RankedTensorType::get(/*shape=*/{}, fp_type);
  auto epsilon_tensor_attr = DenseElementsAttr::get(
      scalar_type, {mlir::cast<Attribute>(epsilon_attr)});
  Value epsilon = b.create<mhlo::ConstantOp>(epsilon_tensor_attr);
  auto dims_type = RankedTensorType::get(/*shape=*/{0}, b.getIntegerType(64));
  auto dims = DenseIntElementsAttr::get(dims_type, SmallVector<int64_t, 1>{});

  if (flatten_tuple) {
    auto operands = llvm::to_vector(values);
    operands.push_back(token);
    return operands;
  }

  auto value = values[0];
  // If value with token already exists, reuse it.
  auto it = rewritten_values.find(value);
  if (it != rewritten_values.end()) return {it->getSecond()};

  auto create_tuple = [&](ArrayRef<Value> operands) {

bool CanBeSymmetricallyQuantized(Value weight) {
  auto dq_op = weight.getDefiningOp<quantfork::DequantizeCastOp>();
  if (!dq_op) return true;

  auto qtype =
      mlir::cast<TensorType>(dq_op.getArg().getType()).getElementType();
  if (auto uniform_type = llvm::dyn_cast_or_null<UniformQuantizedType>(qtype)) {
    return uniform_type.getZeroPoint() == 0;
  } else if (auto per_axis_type =

  Block& function_block = func.front();
  for (auto sharding_and_arg :
       llvm::zip(sharding_for_args, function_block.getArguments())) {
    const auto& sharding = std::get<0>(sharding_and_arg);
    BlockArgument arg = std::get<1>(sharding_and_arg);
    if (failed(VerifySharding(arg.getType(), sharding))) return mlir::failure();
  }
  Operation* terminator = function_block.getTerminator();
  for (auto sharding_and_retval :

// exist.
func::FuncOp GetMainFunction(ModuleOp module_op) {
  const auto main_func_id =
      StringAttr::get(module_op.getContext(), kImportModelDefaultGraphFuncName);
  auto func_ops = module_op.getOps<func::FuncOp>();
  auto main_func_itr = absl::c_find_if(func_ops, [&main_func_id](auto func_op) {
    return func_op.getName() == main_func_id;
  });

  if (main_func_itr == func_ops.end()) return {};

  out = std::get<2>(flattended_labels);

  auto lhs_map_or = EquationToMap(lhs);
  if (!lhs_map_or.has_value()) return std::nullopt;
  auto lhs_map = lhs_map_or.value();

  auto out_map_or = EquationToMap(out);
  if (!out_map_or.has_value()) return std::nullopt;
  auto out_map = out_map_or.value();

  EinsumDimensionNumbers dnums;
  for (int64_t i = 0; i < lhs.size(); ++i) {
    auto out_index = out_map.find(lhs[i]);

  groups_to_only = 0;
  for (const auto& [group, branch] : from) {
    auto to_iter = to.find(group);
    if (to_iter == to.end()) {
      ++groups_from_only;
    } else {
      auto to_branch = to_iter->second;
      if (to_branch == branch) {
        ++groups_same_branch;
      } else {
        ++groups_different_branch;
      }
    }
  }
  for (const auto& [group, _] : to) {

  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  int start = -1, stop = -1;
  auto status = cc_ctx->op_kernel().InputRange(name, &start, &stop);
  args->start = start;
  args->stop = stop;
  tensorflow::Set_TF_Status_from_Status(args->status, status);
}

TF_DataType TF_InputDatatype(TF_OpKernelContext* ctx, int index) {

  if (auto repinput = dyn_cast<TF::TPUReplicatedInputOp>(op)) {
    if (!ValidateReplicatedInput(repinput, num_replicas, tpu_replicate_attr))
      return false;
  }
  if (auto repoutput = dyn_cast<TF::TPUReplicatedOutputOp>(op)) {
    if (!ValidateReplicatedOutput(repoutput, num_replicas, tpu_replicate_attr))
      return false;
  }
  if (auto partinput = dyn_cast<TF::TPUPartitionedInputOp>(op)) {

  SmallVector<int64_t, 2> output_shape{1, 2};
  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
  auto weight_type = RankedTensorType::get(weight_shape, builder->getF32Type());
  auto bias_type = RankedTensorType::get(bias_shape, builder->getF32Type());
  auto projection_type =
      RankedTensorType::get(projection_shape, builder->getF32Type());
  auto layer_norm_scale_type =