// associated first replica operand.
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(input)) {
    if (block_arg.getOwner() != replicate_block) return;

    shape_op.setOperand(replicate_op.GetReplicaOperandForBlockArgument(
        block_arg, /*replica=*/0));

    return;
  }

  Operation* input_def = input.getDefiningOp();

  // If ShapeOp operand is a ReadVariableOp result where the ReadVariableOp

      replicated_inputs.back().push_back(
          replicate.GetReplicaOperandForBlockArgument(arg, i));
    }
    new_replicated_inputs.emplace_back(replicated_inputs.back(), arg.getType());
  }
  for (const auto& arg : replicate.GetPackedBlockArguments()) {
    new_packed_inputs.emplace_back(
        replicate.GetReplicaOperandForBlockArgument(arg, /*replica=*/0));
  }
  SmallVector<Value, 4> new_input_values;

    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();
  }];

  let builders = [

    // 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);

// Returns the operand being forwarded as a replicated/packed block argument for
// a given replica. This assumes a valid block argument (of the replicate op)
// and a valid replica is provided.
Value ReplicateOp::GetReplicaOperandForBlockArgument(BlockArgument block_arg,
                                                     unsigned replica) {
  MutableArrayRef<OpOperand> operands = GetOperandsForBlockArgument(block_arg);