block_arg.getLoc(),
      DenseIntElementsAttr::get(
          RankedTensorType::get({1, num_replicas}, builder.getIntegerType(32)),
          group_assignment_val));

  StringAttr reduce_op = builder.getStringAttr("Add");
  StringAttr mode = builder.getStringAttr("CrossReplica");
  return builder.create<XlaAllReduceOp>(block_arg.getLoc(), block_arg.getType(),

  return op->hasAttr(TF::kReplicationInfoAttr);
}

StringAttr GetReplicationAttr(mlir::Operation* op) {
  return op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
}

StringAttr GetReplicationAttr(TF::TPUCompilationResultOp op) {
  // Special case for getting the replication region for
  // TPUCompilationResultsOp.
  return op->getAttrOfType<StringAttr>(kTpuCompilationStatus);
}

    auto replication_info_attr_str =
        mlir::dyn_cast<StringAttr>(replication_info_attr);
    if (!replication_info_attr_str ||
        replication_info_attr_str.getValue().empty())
      return metadata_op.emitError() << kBadReplicateInfoAttrMsg;

    // Remove `name` attribute.
    attrs.erase(StringAttr::get(metadata_op.getContext(), kNameAttr));

  }

  return {};
}

// Gets the "main" function from the module. Returns an empty op iff it doesn't
// 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) {

                                       value->mutable_func()->mutable_attr()));
  return absl::OkStatus();
}

Status ConvertAttribute(const mlir::StringAttr& attr, AttrValue* value) {
  absl::string_view attr_value(attr.getValue().data(), attr.getValue().size());
  switch (mangling_util::GetMangledKind(attr_value)) {
    case mangling_util::MangledKind::kUnknown: {

}

Value CreateEinsumOpFromXlaDotV2Op(OpBuilder& builder, const Location loc,
                                   Value lhs, Value rhs, Value output,
                                   StringAttr dot_dimension_numbers_str) {
  xla::DotDimensionNumbers dot_dimension_numbers;
  dot_dimension_numbers.ParseFromString(dot_dimension_numbers_str.str());
  SmallVector<Value> input_arguments = {lhs, rhs};

  }

  TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
  return call_op && call_op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
         call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue() ==
             llvm::StringRef(
                 QuantTraitValues[QuantizationTrait::FullyQuantizable]);
}

// Returns the composite function name.

  EXPECT_EQ(
      mlir::cast<StringAttr>(lifted_op->getAttr("_original_entry_function")),
      "composite_dot_general_fn_1");
}

TEST_F(LiftAsFunctionCallTest, EinsumSupportedForXlaDotV2Succeeds) {
  StringAttr einsum_supported_by_xla_dot_v2_attr =
      builder_.getStringAttr("ijk,ikm->ijm");
  StringAttr einsum_one_operand = builder_.getStringAttr("ijk->ikj");

    if (!matchPattern(input, m_Constant(&arg_value))) {
      return failure();
    }
    auto attr_name = signature.getArgAttrOfType<StringAttr>(
        operand.index(), kAttrArgumentNameAttr);
    auto attr_type = signature.getArgAttrOfType<StringAttr>(
        operand.index(), kAttrArgumentTypeAttr);
    auto value = ProcessAttributeValue(arg_value, attr_type);

  // propagated to the reduce function call.
  reduce_call->setAttr(kDeviceAttr,
                       reduce_dataset->getAttrOfType<StringAttr>(kDeviceAttr));
  reduce_call->setAttr(
      TF::kCompileDeviceTypeAttr,
      reduce_dataset->getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr));

  SmallVector<Value, 4> if_returns;

  builder.create<TF::YieldOp>(loc,