code: def test_operators(self, shape, kernel):
      params = kernel.init(jax.random.PRNGKey(0), jnp.zeros(shape))
      lp = kernel.log_prior(params)
      self.assertLess(lp, 0.0)
      output = kernel.apply(params, jnp.ones(shape))
      self.assertEqual(shape[:-1] + (1,), output.shape)
  
      self.assertEqual(
          jax.tree_util.tree_structure(
              kernel.get_all_distributions(),
              is_leaf=lambda x: isinstance(
                  x, distribution_lib.Distribution
              ),
          ),
          jax.tree_util.tree_structure(params["params"]),
      )
  
file_path: probability/tensorflow_probability/python/experimental/autobnn/operators_test.py
imports: from absl.testing import parameterized
  import jax
  import jax.numpy
  import numpy
  from tensorflow_probability.python.experimental.autobnn import kernels
  from tensorflow_probability.python.experimental.autobnn import operators
  from tensorflow_probability.python.experimental.autobnn import util
  from tensorflow_probability.substrates.jax.distributions import distribution
  from absl.testing import absltest
label_desc: Use image node_id 1 for calling the OperatorsTest obj's underlying member method code with example usage: obj.test_operators(shape, kernel) without return types