class_docstr: Add two or more BNNs, with weights taken from a Dirichlet prior.
code: def distributions(self):
      bnn_concentrations = [
          1.0 if isinstance(b, BnnOperator) else 1.5 for b in self.bnns
      ]
      if self.going_to_be_multiplied:
          concentration = self.alpha * jnp.array(bnn_concentrations)
      else:
          concentration = self.alpha * jnp.array(
              [bnn_concentrations for _ in range(self.num_outputs)]
          )
      return super().distributions() | {
          "bnn_weights": dirichlet_lib.Dirichlet(
              concentration=concentration
          )
      }
  
imports: from typing import Optional
  from flax import linen
  import jax.numpy
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.python.experimental.autobnn import likelihoods
  from tensorflow_probability.substrates.jax.bijectors import chain
  from tensorflow_probability.substrates.jax.bijectors import scale
  from tensorflow_probability.substrates.jax.bijectors import shift
  from tensorflow_probability.substrates.jax.distributions import beta
  from tensorflow_probability.substrates.jax.distributions import dirichlet
  from tensorflow_probability.substrates.jax.distributions import half_normal
  from tensorflow_probability.substrates.jax.distributions import normal
  from tensorflow_probability.substrates.jax.distributions import transformed_distribution
label_desc: Use image node_id 1 for calling the WeightedSum obj's underlying member method code with example usage: obj.distributions() and returns: unknown