Easy mechanism to explicitly specify dynamic params

[allen-adastra]

It would be nice if there were a mechanism to state which eqx.Module members are dynamic and a filter transformation that treats only elements marked dynamic as dynamic. This would be helpful for keeping track of what parameters in the model are "trainable" as we build and train larger eqx.Module subclasses.

Thoughts on the idea behind this pattern and how we may achieve it?

One option that comes to mind is to require the user to implement a dynamic_members() method that returns the members that are dynamic.

Another option would be to make a thin wrapper class around dynamic objects.s.