<< replicate.getN() << ")";
  }

  // add block arg to region. This is in $body.
  unsigned block_arg_idx = replicate.GetNumReplicatedBlockArguments();
  Block& block = replicate.GetBody();
  block_arg = block.insertArgument(block_arg_idx, type, replicate.getLoc());

  // add to $replicated_inputs. This also updates OperandSegmentSizeAttr.

  // Build the graph:
  // "add" = "arg0" + "arg1"
  // "ret0" = "add"
  // "ret1" = "arg1"
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_INT32, 0);
  Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_FLOAT, 1);
  Output arg2 = ops::_Arg(s.WithOpName("arg2"), DT_INT32, 2);
  Output add = ops::Add(s.WithOpName("add"), arg0, arg0);

  }
  void getDependentDialects(::mlir::DialectRegistry &registry) const override {}

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    patterns.add<FoldBroadcastInDimBeforeMulOp>(&getContext());
    patterns.add(ConstantFoldMul);
    if (failed(applyPatternsAndFoldGreedily(getOperation(),
                                            std::move(patterns)))) {
      return signalPassFailure();
    }

    // If the parent is not a FuncOp, then add the parent op containing a region
    // to worklist.
    Operation* parent = region->getParentOp();
    if (!isa<func::FuncOp>(parent)) {
      if (ops_with_tokens.insert(parent).second) {
        worklist.push_back(parent);
      }
      continue;
    }

    // For functions, get all the users and add them to the worklist.

  function.setType(builder.getFunctionType(argument_types, result_types));

  // Add resource_name attribute to the output for the resources.
  for (auto& resource : output_only_resources)
    function.setResultAttr(resource.first, kResourceNameArgAttr,
                           builder.getStringAttr(resource.second));

  // Add aliasing_output attribute to the input argument for the resources that
  // are updated by the function.

  MLIRContext* ctx = &getContext();
  auto func = getOperation();

  // The pattern includes decomposing batch normalization ops, fusing add/mul
  // with a constant operand to a preceding affine operation.
  RewritePatternSet patterns(ctx);
  populateWithGenerated(patterns);
  patterns.add<RemoveIdentity, ConstantFoldQuantizableOperands>(ctx);
  if (op_set_ != OpSet::XLA) {
    // Convert Einsum into BatchMatMul for non-XLA opsets.

    rollout_config.enabled_for_gpu_ = true;

    // Set flag to enable using XLA devices. PJRT currently is only supported
    // for XLA devices.
    GetXlaDeviceFlags()->tf_xla_enable_xla_devices = true;

    // Add and setup the GPU device.
    auto device_type = DeviceType(DEVICE_GPU);
    auto jit_device_type = DeviceType(DEVICE_GPU);
    rollout_config.AllowForDeviceInXlaLaunch(device_type);

};

void addPatterns(MLIRContext* context, RewritePatternSet& patterns,
                 bool preserve_assert_op) {
  // Add TF->TF lowering patterns.
  TF::PopulateLoweringTFPatterns(context, &patterns);

  // Add the generated patterns to the list.
  populateWithGenerated(patterns);
  patterns.add<ConvertTFConcatV2Op, ConvertTFBatchMatMulOp,
               ConvertTFBatchMatMulV2Op, ConvertTFBatchMatMulV3Op,

        std::optional<StringRef> composite_function_name =
            GetCompsiteFunctionName(defining_op);
        if (!composite_function_name.has_value()) continue;

        // Do not add CustomAggregatorOp after Gather since it is a weight-only
        // quantizable op.
        if (auto call_op =
                dyn_cast_or_null<TF::PartitionedCallOp>(defining_op)) {
          StringRef function_name =

    return success();
  }
};

// TODO(lyandy): Add canonicalization for empty IslandOps with more than one
// control operand and no data results.

}  // anonymous namespace

void IslandOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
  results.add<DropEmptyIslandNoOperandNoDataResult,