def flatten_weights_and_state(weights, state):
  """Flatten weights and state into lists, excluding empty and cached ones."""
  def _is_empty_weight(x):
    return (x is EMPTY_WEIGHTS or
            (isinstance(x, dict) and x == GET_WEIGHTS_FROM_CACHE))
  flat_weights = [w for w in fastmath.tree_flatten(weights)
                  if not _is_empty_weight(w)]
  def _is_empty_state(x):
    return (x is EMPTY_STATE or
            (isinstance(x, dict) and x == GET_STATE_FROM_CACHE))
  flat_state = [s for s in fastmath.tree_flatten(state)
                if not _is_empty_state(s)]
  return flat_weights, flat_state
Flatten weights and state into lists, excluding empty and cached ones.
trax/layers/base.py