Functional transformations with respect to subsets of parameters

[mjo22]

Hello! I have a question.

I am using equinox for a project here: https://github.com/mjo22/cryojax. My Modules are very nested, and in general I am interested in using grad and vmap with respect to arbitrary subsets of parameters across different Modules (many of my parameters are configured by the user and usually won’t be modified). What is the recommended way to handle this scenario within equinox?

I have some idea of how to handle this, which I have in the README. However, it’s a bit awkward and doesn’t always work. If you have this written somewhere in the docs and I haven’t see it, I apologize!