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