Investigate the use of optree for structured parameter sets
dfm opened this issue · 0 comments
dfm commented
I recently discovered the optree library designed for efficient operations on nested Python data structures (à la JAX Pytrees). It's possible to build a model wrapper that is significantly more flexible than the current named parameters and the performance seems good. This could be fun to play around with! Here's a demo: https://gist.github.com/dfm/9b902c659d728e11e3b48a37aa6990ad
/cc @zachjweiner