code: def test_distribution_methods(self):
      batch_shape = [4]
      sample_shape = [2, 3]
      sample_and_batch_shape = [2, 3, 4]
      d = jit_public_methods.JitPublicMethods(
          normal.Normal(
              loc=0.0, scale=tf.ones(batch_shape, dtype=tf.float32)
          )
      )
      self.assertEqual(d.dtype, tf.float32)
      self.assertAllEqual(d.event_shape, [])
      self.assertAllEqual(d.event_shape_tensor(), [])
      self.assertAllEqual(d.batch_shape, batch_shape)
      self.assertAllEqual(d.batch_shape_tensor(), batch_shape)
  
      if JAX_MODE:
          # TODO(b/184565492): Enable passing shapes to jitted methods in JAX.
          x = tf.zeros(sample_and_batch_shape, dtype=tf.float32)
      else:
          x = d.sample(
              sample_shape,
              seed=test_util.test_seed(sampler_type="stateless"),
          )
      self.assertAllEqual(x.shape, sample_and_batch_shape)
      self.assertAllEqual(d.log_prob(x).shape, sample_and_batch_shape)
  
      self.assertAllEqual(d.mean().shape, batch_shape)
      self.assertAllEqual(d.stddev().shape, batch_shape)
      self.assertAllEqual(d.entropy().shape, batch_shape)
  
      d_copy = d.copy(loc=[1.0, 2.0, 3.0], scale=1.0)
      self.assertIsInstance(d_copy, jit_public_methods.JitPublicMethods)
      self.assertEqual(d_copy.batch_shape, [3])
  
      d_slice = d_copy[:2]
      self.assertIsInstance(
          d_slice, jit_public_methods.JitPublicMethods
      )
      self.assertEqual(d_slice.batch_shape, [2])
  
imports: import types
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.util import jit_public_methods
  from tensorflow_probability.python.internal import test_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 1 for calling the TracedPublicMethodsTest obj's underlying member method code with example usage: obj.test_distribution_methods() without return types