add more detail about aux_data in pytree docs #8739
Conversation
| to specify how to flatten and unflatten instances of the new type. That is, you | ||
| register two functions: the first function accepts instances of the new type and | ||
| returns a pair where the first element is an iterable of children (for example, | ||
| the values of a dictionary) and the second element is auxiliary metadata needed |
There was a problem hiding this comment.
Could we add some words about how you might decide what parts of an object are "children" and what are "aux_data"? My understanding is that children should contain all the traceable/dynamic/JAX-array components (not sure the best language to use there), and aux_data should be all the static components (and must be hashable, I believe)
There was a problem hiding this comment.
Good idea. Though I don't like the word "static", I'll figure out a way to express that.
There was a problem hiding this comment.
(The reason I don't like the word "static" here is that I think of that as jit- or staging-specific jargon, which doesn't cover all cases here, e.g. in how pytree metadata interacts with vmap or AD. It's hard to think of a better term though...)
There was a problem hiding this comment.
Maybe something like "generally all array components should go into children, and all non-array metadata should go into aux_data"
There was a problem hiding this comment.
Nice! That seems like a good rule of thumb.
|
Ping on this - it would be a nice addition to the docs |
fixes #8726