using equinox with xmap

[jakubMitura14]

Hello, Thank you for publishing a fantastic library! Is it possible to use named axes operations (xmap) within init and call of nn .module and then xmap over all created module? I suppose It is as it is just a pytree but I do not have an idea how