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 nested dictionary of distributions for the model's params.
  
  The nested dictionary should have the same structure as the
  variables returned by the init() method, except all leaves should
  be tensorflow probability Distributions.
docstr_tok: ["Return","a","nested","dictionary","of","distributions","for","the","model","'s","params",".","The","nested","dictionary","should","have","the","same","structure","as","the","variables","returned","by","the","init","(",")","method",",","except","all","leaves","should","be","tensorflow","probability","Distributions","."]
code: def distributions(self):
      """Return a nested dictionary of distributions for the model's params.
  
      The nested dictionary should have the same structure as the
      variables returned by the init() method, except all leaves should
      be tensorflow probability Distributions.
      """
      # TODO(thomaswc): Consider having this optionally also be able to
      # return a tfd.JointNamedDistribution, so as to support dependencies
      # between the subdistributions.
      raise NotImplementedError("Subclasses of BNN must define this.")
  
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 1 for calling the BayesianModule obj's underlying member method code with example usage: obj.distributions() without return types