class_docstr: A linen.Module with distributions over its parameters.
  
  Example usage:
    class MyModule(BayesianModule):
  
      def distributions(self):
        return {'dense': {'kernel': tfd.Normal(loc=0, scale=1),
                          'bias': tfd.Normal(loc=0, scale=1)},
                'amplitude': tfd.LogNormal(loc=0, scale=1)}
  
      def setup(self):
        self.dense = nn.Dense(50)
        super().setup()   # <-- Very important, do not forget!
  
      def __call__(self, inputs):
        return self.amplitude * self.dense(inputs)
  
  
    my_bnn = MyModule()
    params = my_bnn.init(jax.random.PRNGKey(0), jnp.zeros(10))
    lp = my_bnn.log_prior(params)
  
  Note that in this example, self.amplitude will be initialized using
  the given tfd.LogNormal distribution, but the self.dense's parameters
  will be initialized using the nn.Dense's default initializers.  However,
  the log_prior score will take into account all of the parameters.
docstr: Children classes must call this from their setup() !
docstr_tok: ["Children","classes","must","call","this","from","their","setup","(",")","!"]
code: def setup(self):
      """Children classes must call this from their setup() !"""
  
      def make_sample_func(dist):
          def sample_func(key, shape):
              return dist.sample(sample_shape=shape, seed=key)
  
          return sample_func
  
      for k, v in self.distributions().items():
          # Create a variable for every distribution that doesn't already
          # have one.  If you define a variable in your setup, we assume
          # you initialize it correctly.
          if not hasattr(self, k):
              try:
                  setattr(
                      self, k, self.param(k, make_sample_func(v), 1)
                  )
              except flax.errors.NameInUseError:
                  # Sometimes subclasses will have parameters where the
                  # parameter name doesn't exactly correspond to the name of
                  # the object field.  This can happen with arrays of parameters
                  # (like PolynomialBBN's hidden parameters.) for example.  I
                  # don't know of any way to detect this beforehand except by
                  # trying to call self.params and having it fail with NameInUseError.
                  # (For example, self.variables doesn't exist at setup() time.)
                  pass
  
imports: import dataclasses
  import flax
  from flax import linen
  import jax.numpy
  from jaxtyping import Array, Float, PyTree
  from tensorflow_probability.python.experimental.autobnn import likelihoods
  from tensorflow_probability.substrates.jax.distributions import distribution
label_desc: Use image node_id 2 for calling the BayesianModule obj's underlying member method code with example usage: obj.setup() and returns: sample_func, dist