// Stores previously seen init types.
  llvm::SmallDenseSet<StringAttr> init_types{};

  for (auto init_sym :
       session_init_op.getInitializers().getAsValueRange<FlatSymbolRefAttr>()) {
    auto init_func = symbol_table.lookup<func::FuncOp>(init_sym);
    if (!init_func) continue;

    auto init_type =
        init_func->getAttrOfType<StringAttr>(kTfSavedModelInitializerTypeAttr);
    if (!init_type) continue;

    attributes.push_back(kLayerNormalizedLstmCellSimple);
  } else {
    attributes.push_back(kLstmCellSimple);
  }

  if (cifg) {
    attributes.push_back(kCoupleInputForgetGates);
  }

  mlir::StringAttr attr_values =
      builder->getStringAttr(llvm::join(attributes, ","));

  func->setAttr(kTFImplements, attr_values);
  return func;
}

class LstmUtilsTest : public ::testing::Test {
 protected:

  std::vector<std::string> disabled_checks;
  for (auto attr : op.getDisabledChecks().getAsRange<StringAttr>()) {
    disabled_checks.push_back(attr.getValue().str());
  }
  std::vector<std::string> platforms;
  for (auto attr : op.getPlatforms().getAsRange<StringAttr>()) {
    platforms.push_back(attr.getValue().str());
  }
  TF_ASSIGN_OR_RETURN(
      auto loader,

  StringAttr interface_name =
      builder.getStringAttr(absl::StrCat("func_", func_count));

  added_func_op->setAttr(kInterfaceNameAttr, interface_name);
  added_call_op->setAttr(kInterfaceNameAttr, interface_name);

  StringAttr device = mlir::cast<StringAttr>(
      added_func_op->getRegion(0).getBlocks().front().front().getAttr(kDevice));
  StringAttr inference_type = mlir::cast<StringAttr>(

  using DeviceNameUtils = tensorflow::DeviceNameUtils;

  // Duplicate of the logic in MultiDeviceFunctionBodyPlacer::BodyNodeDevice
  // LINT.IfChange
  auto device_id = StringAttr::get(call->getContext(), "device");
  auto caller_device = call->getAttrOfType<StringAttr>(device_id);
  if (!caller_device) return;

  DeviceNameUtils::ParsedName caller_parsed_device;
  if (!DeviceNameUtils::ParseFullName(caller_device.getValue().str(),

    if (!is_generator(alias)) return true;

    StringAttr device;
    if (auto arg = mlir::dyn_cast<BlockArgument>(alias)) {
      device = func.getArgAttrOfType<mlir::StringAttr>(arg.getArgNumber(),
                                                       kFuncDeviceAttr);
    } else {
      device = alias.getDefiningOp()->getAttrOfType<StringAttr>(kDeviceAttr);
    }

    if (!device) return false;

  return std::make_pair(f_attr.getValue().str(), kEmptyNodeName.str());
}

Operation *DuplicateOp(TF::PartitionedCallOp call_op, PatternRewriter &rewriter,
                       const StringAttr &new_ref_func_name) {
  // Create PartitionedCallOp to the copied composite function. This
  // PartitionedCallOp does not have kQuantTraitAttrName, and therefore won't
  // get quantized.

  if (!mlir::isa<mlir::StringAttr>(shard_attr)) return mlir::failure();

  auto shard_str = mlir::cast<mlir::StringAttr>(shard_attr).getValue().str();
  return DecodeShardingAttribute(shard_str, sharding, report_error);
}

void EncodeSharding(mlir::Operation* op, llvm::StringRef shard_str) {
  if (!op->hasAttrOfType<mlir::StringAttr>(shard_str)) return;

  ::xla::OpSharding sharding;

  if (dict_attr) {
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
        << "inputs missing in entry function attribute";
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
        << "outputs missing in entry function attribute";
    mlir::cast<mlir::StringAttr>(dict_attr.get("inputs"))
        .getValue()

  if (dict_attr) {
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
        << "inputs missing in entry function attribute";
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
        << "outputs missing in entry function attribute";
    mlir::cast<mlir::StringAttr>(dict_attr.get("inputs"))
        .getValue()