class_docstr: A BNN where samples are polynomial functions.
class_docstr_tok: ["A","BNN","where","samples","are","polynomial","functions","."]
code: def distributions(self):
      d = super().distributions()
      del d["dense1"]
      for i in range(self.degree):
          # Do not scale these layers by 1/sqrt(width), because we also
          # multiply these weights by the learned `amplitude` parameter.
          d[f"hiddens_{i}"] = {
              "kernel": normal_lib.Normal(loc=0, scale=1.0),
              "bias": normal_lib.Normal(loc=0, scale=self.bias_scale),
          }
      return d | {
          "shift": normal_lib.Normal(
              loc=self.shift_mean, scale=self.shift_scale
          ),
          "amplitude": lognormal_lib.LogNormal(
              loc=0, scale=self.amplitude_scale
          ),
      }
  
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 1 for calling the PolynomialBNN obj's underlying member method code with example usage: obj.distributions() and returns: unknown