code: def test_random_tree(self):
      r0 = bnn_tree.random_tree(
          jax.random.PRNGKey(0), depth=0, width=50, period=7
      )
      self.assertIsInstance(r0, kernels.OneLayerBNN)
      params = r0.init(jax.random.PRNGKey(1), jnp.zeros(5))
      lp = r0.log_prior(params)
      self.assertLess(lp, 0.0)
      output = r0.apply(params, jnp.ones(5))
      self.assertEqual((1,), output.shape)
  
      r1 = bnn_tree.random_tree(
          jax.random.PRNGKey(0), depth=1, width=50, period=24
      )
      self.assertIsInstance(r1, nn.Module)
      params = r1.init(jax.random.PRNGKey(1), jnp.zeros(5))
      lp = r1.log_prior(params)
      self.assertLess(lp, 0.0)
      output = r1.apply(params, jnp.ones(5))
      self.assertEqual((1,), output.shape)
  
      r2 = bnn_tree.random_tree(
          jax.random.PRNGKey(0), depth=2, width=50, period=52
      )
      self.assertIsInstance(r2, nn.Module)
      params = r2.init(jax.random.PRNGKey(1), jnp.zeros(5))
      lp = r2.log_prior(params)
      self.assertLess(lp, 0.0)
      output = r2.apply(params, jnp.ones(5))
      self.assertEqual((1,), output.shape)
  
file_path: probability/tensorflow_probability/python/experimental/autobnn/bnn_tree_test.py
imports: from absl.testing import parameterized
  from flax import linen
  import jax
  import jax.numpy
  from tensorflow_probability.python.experimental.autobnn import bnn_tree
  from tensorflow_probability.python.experimental.autobnn import kernels
  from absl.testing import absltest
label_desc: Use image node_id 3 for calling the TreeTest obj's underlying member method code with example usage: obj.test_random_tree() without return types