def _match_by_shape(full, partial):
  """Puts partial into full matching by shape."""
  partial_idx = 0
  res = []
  for w in full:
    if partial_idx >= len(partial):
      res.append(w)  # read everything from parial list, just fill
    elif w is None and partial[partial_idx] is None:  # both Nones, move on
      res.append(None)
      partial_idx += 1
    elif w is None or partial[partial_idx] is None:  # one None but not both
      res.append(w)
    elif w.shape == partial[partial_idx].shape:
      res.append(partial[partial_idx])
      partial_idx += 1
    else:
      res.append(w)
  if partial_idx < len(partial):
    _log('Did not manage to match shapes in model for all checkpoint weights.')
    for w in partial[:partial_idx]:
      _log('  Inserted tensor of shape %s' % str(w.shape))
    for i, w in enumerate(partial[partial_idx:]):
      _log('  Not inserted tensor of shape %s' % str(w.shape))
      model_weight_shape = str(full[i + partial_idx].shape)
      _log('  Tensor in that place has shape: %s' % model_weight_shape)
    raise IndexError
  return res
Puts partial into full matching by shape.
trax/supervised/training.py