code: def test_methods_are_wrapped(self):
      d = jit_public_methods.JitPublicMethods(
          normal.Normal(0.0, 1.0),
          methods_to_exclude=(
              "event_shape_tensor",
              "batch_shape_tensor",
          ),
      )
      # Wrapped methods should be some sort of non-method type, e.g.:
      #   tensorflow.python.eager.def_function.Function (TF)
      #   types.FunctionType (JAX)
      self.assertNotIsInstance(d.sample, types.MethodType)
      self.assertNotIsInstance(d.log_prob, types.MethodType)
  
      # Excluded methods should not be wrapped.
      self.assertIsInstance(d.event_shape_tensor, types.MethodType)
      self.assertIsInstance(d.batch_shape_tensor, types.MethodType)
  
      # Private methods should not be wrapped.
      self.assertIsInstance(d._sample_n, types.MethodType)
      self.assertIsInstance(d._log_prob, types.MethodType)
  
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 2 for calling the TracedPublicMethodsTest obj's underlying member method code with example usage: obj.test_methods_are_wrapped() without return types