class_docstr: A BNN where samples are polynomial functions.
class_docstr_tok: ["A","BNN","where","samples","are","polynomial","functions","."]
code: def setup(self):
      kernel_init = nn.initializers.normal(1.0)
  
      def bias_init(seed, shape, dtype=jnp.float32):
          return self.bias_init_amplitude * jax.random.normal(
              seed, shape, dtype=dtype
          )
  
      self.hiddens = [
          nn.Dense(
              self.width, kernel_init=kernel_init, bias_init=bias_init
          )
          for _ in range(self.degree)
      ]
      super().setup()
  
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 PolynomialBNN obj's underlying member method code with example usage: obj.setup() and returns: unknown