Bayesian NNs with Equinox

[cdagher]

This is really an exploratory question about implementing a Bayesian layer in Equinox. My research requires Bayesian neural networks, and I've been playing with Equinox to see how it stacks up to Flax - definitely impressed with the ease of using Equinox.

I'm thinking that it shouldn't be too hard to add this functionality with another Module in nn , but I figured I'd throw the question out there and see if there are any other thoughts or suggestions. Maybe a library like NumPyro could be used alongside Equinox?