flax

9ccd27d5 - extended pytreenode to take dataclass kwargs