This repository was archived by the owner on May 11, 2023. It is now read-only.
This repository was archived by the owner on May 11, 2023. It is now read-only.
feat: Use PyTreeNode inplace of PyTree #19
Description
Feature Request
Right now, a custom PyTree object is used. This should be replaced with Flax's PyTreeNode. This would then be consumed by the Parameters object to enable more robust parameter handling.