code: def flatten(weights, start=0, stop=2):
      
      for key, val in weights.items():
          new_shape = val.shape[0:start] + (-1, ) + val.shape[stop:]
          weights[key] = val.reshape(new_shape)
      return weights
docs: This methods reshapes all values in a dictionary.
  
      The indices from start to stop will be flattened into a single index.
  
      Args:
          weights: A dictionary mapping keys to numpy arrays.
          start: The starting index.
          stop: The ending index.