"math"
	"math/bits"
)

// checkFunc checks invariants of f.
func checkFunc(f *Func) {
	blockMark := make([]bool, f.NumBlocks())
	valueMark := make([]bool, f.NumValues())

	for _, b := range f.Blocks {
		if blockMark[b.ID] {
			f.Fatalf("block %s appears twice in %s!", b, f.Name)
		}
		blockMark[b.ID] = true
		if b.Func != f {
			f.Fatalf("%s.Func=%s, want %s", b, b.Func.Name, f.Name)
		}

  cast_input.setType(transform_result_type);
  auto block_arg = mlir::dyn_cast<mlir::BlockArgument>(cast_input);
  auto cast_op_input = dyn_cast_or_null<TF::CastOp>(cast_input.getDefiningOp());
  while (block_arg || cast_op_input) {
    if (block_arg) {
      // Change on device function type/shape.
      HandleFuncOp(block_arg.getOwner()->getParentOp());
      block_arg = nullptr;
      cast_op_input = nullptr;
    } else {

             << ") and block argument " << block_arg.getArgNumber() << " ("
             << block_arg.getType() << ") to have compatible types";

    return success();
  };
  for (auto block_arg : block.getArguments()) {
    if (block_arg.getArgNumber() < num_replicated_block_args) {
      for (int32_t i = n * block_arg.getArgNumber(), e = i + n; i < e; ++i)
        if (failed(verify_operand_types(block_arg, i))) return failure();
    } else {

  }

  BlockArgument block_arg;
  if (failed(AppendReplicatedInput(builder, replicate, inputs,
                                   val_bcast.getType(), block_arg))) {
    return failure();
  }

  OpBuilder before_cluster_builder(cluster);
  IdentityOp assigned_id = before_cluster_builder.create<IdentityOp>(
      val_bcast.getLoc(), block_arg.getType(), block_arg);

      continue;
    }
    auto block_arg = mlir::dyn_cast<BlockArgument>(arg);
    if (!block_arg || block_arg.getOwner() != &replicate.GetBody()) continue;
    assert(replicate_arg_to_execute_arg.count(block_arg.getArgNumber()) == 0 &&
           "Found duplicate use of a resource in the execute op.");
    replicate_arg_to_execute_arg[block_arg.getArgNumber()] =
        index_and_arg.index();
  }

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

      const int64_t execute_arg_index = input_and_idx.index();
      if (auto block_arg = mlir::dyn_cast<BlockArgument>(input)) {
        // For a block argument, consider transforms only when it is a
        // replicated input (defining ops will be outside the replicate node).
        if (maybe_replicate != block_arg.getParentRegion()->getParentOp() ||
            !HandleReplicatedInputs(execute_arg_index, execute.getKey(),

      auto block_arg =
          mlir::dyn_cast<BlockArgument>(ret->getOperand(result_idx));
      if (!block_arg) {
        return op->emitOpError("result #")
               << result_idx << " not tied to function argument for branch @"
               << func.getName();
      }
      if (!common_arg_index.has_value()) {
        common_arg_index = block_arg.getArgNumber();

        island_op.getLoc(), output_types, control_type, replica_inputs);

    // Map block arg to replica arg.
    mapping.clear();
    for (auto& block_arg : replicate_op.GetBody().getArguments())
      mapping.map(block_arg,
                  replicate_op.GetReplicaOperandForBlockArgument(block_arg, i));

    // Copy over replicate region into replica island.
    replicate_op.getBody().cloneInto(&replica.getBody(), mapping);

      // If the value is defined as an argument of the func_op, adds it to
      // the argument list of the function that uses this op.
      if (BlockArgument block_arg = mlir::dyn_cast<BlockArgument>(value)) {
        if (StringAttr attr = func_op.getArgAttrOfType<StringAttr>(
                block_arg.getArgNumber(), kTFDeviceAttr)) {
          value_device = attr.getValue().str();
        }

        if (GetHost(value_device) != op_host) {