class_docstr: A BNN with one hidden layer.
class_docstr_tok: ["A","BNN","with","one","hidden","layer","."]
code: def distributions(self):
      # Strictly speaking, these distributions don't exactly correspond to
      # the initializations used in setup().  lecun_normal uses a truncated
      # normal, for example, and the zeros_init used for the bias certainly
      # isn't a sample from a normal.
      d = {
          "dense1": {
              "kernel": normal_lib.Normal(
                  loc=0, scale=1.0 / jnp.sqrt(self.width)
              ),
              "bias": normal_lib.Normal(loc=0, scale=self.bias_scale),
          }
      }
      if not self.going_to_be_multiplied:
          d["dense2"] = {
              "kernel": normal_lib.Normal(
                  loc=0, scale=1.0 / jnp.sqrt(self.width)
              ),
              "bias": normal_lib.Normal(loc=0, scale=self.bias_scale),
          }
      return super().distributions() | d
  
imports: from flax import linen
  from flax.linen import initializers
  import jax
  import jax.numpy
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.substrates.jax.distributions import lognormal
  from tensorflow_probability.substrates.jax.distributions import normal
  from tensorflow_probability.substrates.jax.distributions import student_t
  from tensorflow_probability.substrates.jax.distributions import uniform
label_desc: Use image node_id 2 for calling the OneLayerBNN obj's underlying member method code with example usage: obj.distributions() and returns: unknown