code: def _param_func_generator(
      self,
      data,
      dist_name,
      params,
      batch_shape,
      func,
      generate_sample_function=False,
  ):
      for param_name, param in params.items():
          if not tf.is_tensor(param) or not np.issubdtype(
              param.dtype, np.floating
          ):
              continue
  
          def _dist_func(param_name, param):
              return data.draw(
                  self._make_distribution(
                      dist_name,
                      params,
                      batch_shape,
                      override_params={param_name: param},
                  )
              )
  
          def _func(param_name, param):
              return func(_dist_func(param_name, param))
  
          yield param_name, param, _dist_func, _func
  
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 2 for calling the _GradTest obj's underlying member method code with example usage: obj._param_func_generator(data, dist_name, params, batch_shape, func, generate_sample_function) and returns: data, func