code: def random_pytree(key, treedef):
      keys = jr.split(key, treedef.num_leaves)
      leaves = []
      for key in keys:
          dimkey, sizekey, valuekey = jr.split(key, 3)
          num_dims = jr.randint(dimkey, (), 0, 5).item()
          dim_sizes = jr.randint(sizekey, (num_dims,), 0, 5).tolist()
          value = jr.normal(valuekey, dim_sizes)
          leaves.append(value)
      return jtu.tree_unflatten(treedef, leaves)
  
imports: import equinox
  import jax.random
  import jax.tree_util
label_desc: Use image node_id 1 for calling a global function with example usage: random_pytree(key, treedef) and returns: jtu