code: def test_weighted_sum_of_all(self, depth):
      soa = bnn_tree.weighted_sum_of_all(
          jnp.linspace(0.0, 1.0, 100), jnp.ones(100), depth=depth
      )
      params = soa.init(jax.random.PRNGKey(0), jnp.zeros(5))
      lp = soa.log_prior(params)
      self.assertLess(lp, 0.0)
      output = soa.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 2 for calling the TreeTest obj's underlying member method code with example usage: obj.test_weighted_sum_of_all(depth) without return types