def unflatten_weights_and_state(
    flat_weights, flat_state, weights_and_state_signature, weights_only=False):
  """Unflatten weights and state given their signatures."""
  weights_tree, state_tree = weights_and_state_signature
  weights_to_copy = [EMPTY_WEIGHTS, GET_WEIGHTS_FROM_CACHE]
  weights, _ = fastmath.tree_unflatten(flat_weights, weights_tree,
                                       copy_from_tree=weights_to_copy)
  state = None
  if not weights_only:
    states_to_copy = [EMPTY_STATE, GET_STATE_FROM_CACHE]
    state, _ = fastmath.tree_unflatten(flat_state, state_tree,
                                       copy_from_tree=states_to_copy)
  return weights, state
Unflatten weights and state given their signatures.
trax/layers/base.py