code: def testDistribution(self):
      def f(s):
          return tfd.Normal(
              loc=np.arange(3.0).astype(np.float32), scale=s
          ).sample(seed=jax.random.PRNGKey(1))
  
      vecf = jax.vmap(f, in_axes=(0,))
      f(np.array(1.0, np.float32))
      f(np.array([1, 2, 3.0], np.float32))
      vecf(np.linspace(1, 3, 10).astype(np.float32))
      vecf(np.linspace(1, 3, 30).reshape(10, 3).astype(np.float32))
  
      jax.jacfwd(f)(np.array(1.0, np.float32))
      jax.jacfwd(f)(np.array([1, 2, 3.0], np.float32))
      jax.jacfwd(vecf)(np.linspace(1, 3, 10).astype(np.float32))
  
file_path: probability/tensorflow_probability/substrates/meta/jax_integration_test.py
imports: import jax
  import jax.numpy
  from tensorflow_probability.substrates import jax
  from absl.testing import absltest
label_desc: Use image node_id 2 for calling the JaxIntegrationTest obj's underlying member method code with example usage: obj.testDistribution() and returns: tfd