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 set_likelihood_model(
      self, likelihood_model: likelihoods.LikelihoodModel
  ):
      self.likelihood_model = likelihood_model
  
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 BNN obj's underlying member method code with example usage: obj.set_likelihood_model(likelihood_model) without return types