code: def test_mybnn(self):
      my_bnn = MyBNN()
      d = my_bnn.distributions()
      self.assertIn("noise_scale", d)
      sample_noise = d["noise_scale"].sample(
          1, seed=jax.random.PRNGKey(0)
      )
      self.assertEqual((1,), sample_noise.shape)
  
      params = my_bnn.init(jax.random.PRNGKey(0), jnp.zeros(1))
      lp1 = my_bnn.log_prior(params)
      params["params"]["amplitude"] += 50
      lp2 = my_bnn.log_prior(params)
      self.assertLess(lp2, lp1)
  
      data = jnp.array(
          [[0], [1], [2], [3], [4], [5]], dtype=jnp.float32
      )
      obs = jnp.array([1, 0, 1, 0, 1, 0], dtype=jnp.float32)
      ll = my_bnn.log_likelihood(params, data, obs)
      lp = my_bnn.log_prob(params, data, obs)
      self.assertLess(jnp.sum(lp), jnp.sum(ll))
  
file_path: probability/tensorflow_probability/python/experimental/autobnn/bnn_test.py
imports: from flax import linen
  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 absl.testing import absltest
label_desc: Use image node_id 1 for calling the BnnTests obj's underlying member method code with example usage: obj.test_mybnn() without return types