code: def random_tree(
      key: jax.Array,
      depth: int,
      width: int,
      period: float,
      parent_is_multiply: bool = False,
  ) -> nn.Module:
      """Return a random complete tree BNN of the given depth.
  
      Args:
        key: Random number key.
        depth: Return a BNN of this tree depth.  Zero based, so depth=0 returns
          a leaf BNN.
        width: The number of hidden nodes in the leaf layers.
        period: The period of any PeriodicBNN kernels in the tree.
        parent_is_multiply: If true, don't create a weight layer after the hidden
          nodes of any leaf kernels and only use addition as an internal node.
  
      Returns:
        A BNN of the specified tree depth.
      """
      if depth == 0:
          c = jax.random.choice(key, len(LEAVES))
          return LEAVES[c](
              width=width,
              going_to_be_multiplied=parent_is_multiply,
              period=period,
          )
  
      key1, key2, key3 = jax.random.split(key, 3)
      if parent_is_multiply:
          c = 1  # Can't multiply Multiply or ChangePoints
          is_multiply = True
      else:
          c = jax.random.choice(key1, len(OPERATORS))
          is_multiply = c == 0
  
      sub1 = random_tree(key2, depth - 1, width, period, is_multiply)
      sub2 = random_tree(key3, depth - 1, width, period, is_multiply)
  
      return OPERATORS[c](bnns=(sub1, sub2))
  
file_path: probability/tensorflow_probability/python/experimental/autobnn/bnn_tree.py
imports: from typing import Iterable, List
  from flax import linen
  import jax
  import jax.numpy
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.python.experimental.autobnn import kernels
  from tensorflow_probability.python.experimental.autobnn import operators
  from tensorflow_probability.python.experimental.autobnn import util
label_desc: Use image node_id 3 for calling a global function with example usage: random_tree(key, depth, width, period, parent_is_multiply) and returns: OPERATORS, LEAVES