}

// Updates uses of `old_read` to `new_partitioned_input` and `new_reads`.
// `old_partitioned_input` is the predecessor of `old_read`. `new_reads`
// contains the predecessors of `new_partitioned_input`.
LogicalResult UpdateReadUses(TF::ReadVariableOp old_read,
                             TF::TPUPartitionedInputV2Op old_partitioned_input,
                             TF::TPUPartitionedInputV2Op new_partitioned_input,