xla::XlaComputation* xla_computation, bool use_tuple_args,
    bool enable_op_fallback, bool return_tuple,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns =
        {},
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes = {},
    llvm::StringRef module_name = llvm::StringRef());

// Creates a MLIR pipeline that lowers MLIR module to MHLO dialect. The input

// true then the pipeline will include all the legalization passes.
void CreateConvertMlirToXlaHloPipeline(
    mlir::OpPassManager& pm, llvm::StringRef device_type,
    bool enable_op_fallback,
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes,
    bool lower_to_xla_hlo, bool allow_partial_conversion) {
  bool legalize_chlo = true;

  pm.addNestedPass<mlir::func::FuncOp>(

    const TF::ResourceAliasAnalysis::Info& resource_alias_analysis) {
  // We need to know the devices to copy to.
  if (!replicate.getDevices()) return false;

  MutableArrayRef<OpOperand> inputs =
      replicate.GetOperandsForBlockArgument(replicate_arg);
  for (const auto& entry : llvm::enumerate(inputs)) {
    auto input_op = entry.value().get().getDefiningOp();

//
// Requires that block argument is of this replicate op.
MutableArrayRef<OpOperand> ReplicateOp::GetOperandsForBlockArgument(
    BlockArgument block_arg) {
  assert(block_arg.getOwner() == &GetBody());

  unsigned arg_number = block_arg.getArgNumber();
  unsigned num_replicated_args = GetNumReplicatedBlockArguments();
  int32_t num_replicas = getNAttr().getInt();
  MutableArrayRef<OpOperand> operands = getOperation()->getOpOperands();

Status CompileMlirToXlaHloViaBuilder(
    mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
    llvm::StringRef device_type, XlaCompilationResult* compilation_result,
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes) {
  // This call to RefineShapes is redundant with the call in BuildHloFromTf.
  // It's here so xla::Parameters that are created form block.getArguments will

    unsigned GetReplicaOperandIndexForBlockArgument(BlockArgument block_arg, unsigned replica);
    Value GetReplicaOperandForBlockArgument(BlockArgument block_arg, unsigned replica);
    MutableArrayRef<OpOperand> GetOperandsForBlockArgument(BlockArgument block_arg);
    bool WrapsSingleOp();
  }];

  let builders = [
    OpBuilder<(ins "int":$n,

  return new_island;
}

// Creates respective YieldOp for the new merged island.
YieldOp CreateNewIslandYieldOp(IslandOp new_island,
                               llvm::MutableArrayRef<IslandResult> results) {
  llvm::SmallVector<Value, 8> yield_operands;
  yield_operands.reserve(results.size());

  for (auto [old_result, new_island] :
       llvm::zip(results, new_island.getOutputs())) {

  return mapping;
}

// Adds a new replicated input to the replicate op.
tf_device::ReplicateOp AddInputsToReplicateOp(
    tf_device::ReplicateOp replicate,
    MutableArrayRef<TF::VarHandleOp> new_inputs,
    const llvm::SmallDenseMap<llvm::StringRef, llvm::SmallVector<StringRef, 4>>&
        devices) {
  int64_t num_replicas = replicate.getN();
  assert(new_inputs.size() == num_replicas);

                                tf_device::ReplicateOp replicate,
                                int32_t block_size) {
  // We need to know the devices to copy to.
  if (!replicate.getDevices()) return false;

  MutableArrayRef<OpOperand> inputs =
      replicate.GetOperandsForBlockArgument(block_arg);
  for (auto& input : inputs) {
    auto input_op = input.get().getDefiningOp();

    LLVM_DEBUG(llvm::dbgs() << GemmStyleOp::getOperationName()
                            << " op must have ranked tensor type.\n");
    return failure();
  }

  MutableArrayRef<BlockArgument> operands =
      entry_func_op.getBody().getArguments();
  // Function must have input, filter, and optionally bias.
  if (operands.size() != 2 && operands.size() != 3) {