(  # pylint: disable=invalid-name
    start, end, from_trainer, to_trainer, copy_optimizer_slots=False
)
Copy model weights[start:end] from from_trainer to to_trainer.
Copy model weights[start:end] from from_trainer to to_trainer.
Copy model weights [ start : end ] from from_trainer to to_trainer .
def _copy_model_weights_and_state(  # pylint: disable=invalid-name
    start, end, from_trainer, to_trainer, copy_optimizer_slots=False
):
  """Copy model weights[start:end] from from_trainer to to_trainer."""
  from_weights = from_trainer.model_weights
  to_weights = list(to_trainer.model_weights)
  shared_weights = from_weights[start:end]
  to_weights[start:end] = shared_weights
  to_trainer.model_weights = to_weights

  from_state = from_trainer.model_state
  to_state = list(to_trainer.model_state)
  shared_state = from_state[start:end]
  to_state[start:end] = shared_state
  to_trainer.model_state = to_state

  if copy_optimizer_slots:
    # TODO(lukaszkaiser): make a nicer API in Trainer to support this.
    # Currently we use the hack below. Note [0] since that's the model w/o loss.
    # pylint: disable=protected-access
    from_slots = from_trainer._opt_state.slots[0][start:end]
    to_slots = to_trainer._opt_state.slots[0]
    # The lines below do to_slots[start:end] = from_slots, but on tuples.
    new_slots = to_slots[:start] + from_slots[start:end] + to_slots[end:]
    new_slots = tuple([new_slots] + list(to_trainer._opt_state.slots[1:]))
    to_trainer._opt_state = to_trainer._opt_state._replace(slots=new_slots)
def _copy_model_weights_and_state ( # pylint: disable=invalid-name start , end , from_trainer , to_trainer , copy_optimizer_slots = False ) : from_weights = from_trainer . model_weights to_weights = list ( to_trainer . model_weights ) shared_weights = from_weights [ start : end ] to_weights [ start : end ] = shared_weights to_trainer . model_weights = to_weights from_state = from_trainer . model_state to_state = list ( to_trainer . model_state ) shared_state = from_state [ start : end ] to_state [ start : end ] = shared_state to_trainer . model_state = to_state if copy_optimizer_slots : # TODO(lukaszkaiser): make a nicer API in Trainer to support this. # Currently we use the hack below. Note [0] since that's the model w/o loss. # pylint: disable=protected-access from_slots = from_trainer . _opt_state . slots [ 0 ] [ start : end ] to_slots = to_trainer . _opt_state . slots [ 0 ] # The lines below do to_slots[start:end] = from_slots, but on tuples. new_slots = to_slots [ : start ] + from_slots [ start : end ] + to_slots [ end : ] new_slots = tuple ( [ new_slots ] + list ( to_trainer . _opt_state . slots [ 1 : ] ) ) to_trainer . _opt_state = to_trainer . _opt_state . _replace ( slots = new_slots )