Add Restructure

[gvcallen]

Adds the Restructure bijector similar to tfb.Restructure.

@lockwo I'm not sure how you feel about adding these types of bijectors into distreqx, however I feel they are very necessary if the end goal is a self-contained package with no interop with Tensorflow. Although I'm not sure which direction you are intending for the library to go, I feel there's a definite need in the Equinox ecosystem for a native distribution + bijector library, which these kinds of classes are necessary for.

As an example "critical" use case for this Restructure bijector, I have found it is essentially impossible to easily model correlations between arbitrary pytree leaves without it, without maintaining lots of boilerplate. In my specific case, I am training a normalizing flow on samples of specific leaves in my tree. Restructure can then be used (along with Split) to restructure the flow's output so that it properly represents a distribution over my specific tree, as opposed to only a sub-portion of my tree's leaves flattened which then I must map manually. Combined with a Joint distribution and potentially a multi-bijector similar to tbf's JointMap, and I think its then possible to model arbitrary correlations across arbitrary pytree structures and leaves.

[lockwo]

Yea I think these are in general ok add. Although distrax doesn't have them, I think distrax was very array focused. I think most users (including myself) generally work with arrays, and this pytree stuff is strictly additive so users concerned with only arrays are ok either way.

[gvcallen]

@lockwo ah I see, yes I'm personally dealing with very nested Pytrees all the time :D