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: Return a string summarizing the structure of the BNN.
docstr_tok: ["Return","a","string","summarizing","the","structure","of","the","BNN","."]
code: def summarize(self, params=None, full: bool = False) -> str:
      """Return a string summarizing the structure of the BNN."""
      return self.shortname()
  
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 5 for calling the BayesianModule obj's underlying member method code with example usage: obj.summarize(params, full) and returns: self