Orthogonal reparameterization for Dense layer in Flax/Jax

[Chutlhu]

Dear all,
is there an equivalent of the following pytorch reparameterization in Flax?
https://pytorch.org/docs/stable/generated/torch.nn.utils.parametrizations.orthogonal.html

[davisyoshida]

This would be pretty easy to implement using https://github.com/davisyoshida/qax, a tool I made to make writing this class of transformation easier.

[chiamp]

Flax does not have a specific orthogonal parametrization function, but we have nn.map_variables, which is the Flax equivalent of PyTorch's parametrize