class_docstr: A Bayesian Neural Network.
  
  A BNN's __call__ method must accept a tensor of shape (..., num_features)
  and return a tensor of shape (..., likelihood_model.num_outputs()).
  Given that, it provides log_likelihood and log_prob methods based
  on the provided likelihood_model.
code: def distributions(self):
      # Children classes must call super().distributions() to include this!
      return self.likelihood_model.distributions()
  
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 BNN obj's underlying member method code with example usage: obj.distributions() and returns: self