code: def test_list_of_all(self):
      l0 = bnn_tree.list_of_all(jnp.linspace(0.0, 100.0, 100), 0)
      # With no periods, there should be five kernels.
      self.assertLen(l0, 5)
      for k in l0:
          self.assertFalse(k.going_to_be_multiplied)
  
      l0 = bnn_tree.list_of_all(
          100, 0, 50, [20.0, 40.0], parent_is_multiply=True
      )
      self.assertLen(l0, 7)
      for k in l0:
          self.assertTrue(k.going_to_be_multiplied)
  
      l1 = bnn_tree.list_of_all(jnp.linspace(0.0, 100.0, 100), 1)
      # With no periods, there should be
      # 15 trees with a Multiply top node,
      # 15 trees with a WeightedSum top node, and
      # 25 trees with a LearnableChangePoint top node.
      self.assertLen(l1, 55)
  
      # Check that all of the BNNs in the tree can be trained.
      for k in l1:
          params = k.init(jax.random.PRNGKey(0), jnp.zeros(5))
          lp = k.log_prior(params)
          self.assertLess(lp, 0.0)
          output = k.apply(params, jnp.ones(5))
          self.assertEqual((1,), output.shape)
  
      l1 = bnn_tree.list_of_all(
          jnp.linspace(0.0, 100.0, 100),
          1,
          50,
          [20.0, 40.0],
          parent_is_multiply=True,
      )
      # With 2 periods and parent_is_multiply, there are only WeightedSum top
      # nodes, with 7*8/2 = 28 trees.
      self.assertLen(l1, 28)
  
      l2 = bnn_tree.list_of_all(jnp.linspace(0.0, 100.0, 100), 2)
      # With no periods, there should be
      # 15*16/2 = 120 trees with a Multiply top node,
      # 55*56/2 = 1540 trees with a WeightedSum top node, and
      # 55*55 = 3025 trees with a LearnableChangePoint top node.
      self.assertLen(l2, 4685)
  
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 1 for calling the TreeTest obj's underlying member method code with example usage: obj.test_list_of_all() without return types