Block& block = replicate.GetBody();
  block_arg = block.insertArgument(block_arg_idx, type, replicate.getLoc());

  // add to $replicated_inputs. This also updates OperandSegmentSizeAttr.
  replicate.getReplicatedInputsMutable().append(inputs);

  return success();
}

// Insert an XlaAllReduce.
Value CreateAllReduce(ReplicateOp replicate, OpBuilder& builder,
                      Value block_arg) {