code: def _test_transformation(self, func, param, msg=None):
      primal, tangent = jax.jvp(func, (param,), (np.ones_like(param),))
      self.assertEqual(primal.shape, tangent.shape)
      if not FLAGS.execute_only:
          self.assertNotAllEqual(
              tangent, np.zeros_like(tangent), msg=msg
          )
  
imports: import functools
  import os
  from absl import flags
  from absl.testing import parameterized
  import hypothesis
  from hypothesis import strategies
  import jax
  from jax import random
  import jax.numpy
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal.backend import jax
  from tensorflow_probability.substrates.jax.bijectors import bijector_test_util
  from tensorflow_probability.substrates.jax.distributions import hypothesis_testlib
  from tensorflow_probability.substrates.jax.distributions import normal
  from tensorflow_probability.substrates.jax.distributions import transformed_distribution
  from tensorflow_probability.substrates.jax.internal import hypothesis_testlib
  from tensorflow_probability.substrates.jax.internal import test_util
label_desc: Use image node_id 1 for calling the JVPTest obj's underlying member method code with example usage: obj._test_transformation(func, param, msg) without return types