code: def augment_transition_fn_with_parameters(
      parameter_prior,
      parameterized_transition_fn,
      parameter_constraining_bijector,
  ):
      """Wraps a transition fn on states to act on `ParametersAndState` tuples."""
  
      def params_and_state_transition_fn(
          step, params_and_state, perturbation_scale, **kwargs
      ):
          """Transition function operating on a `ParamsAndState` namedtuple."""
          # Extract the state, to pass through to the observation fn.
          unconstrained_params, state = params_and_state
          if "state_history" in kwargs:
              kwargs["state_history"] = kwargs["state_history"].state
  
          # Perturb each (unconstrained) parameter with normally-distributed noise.
          if not tf.nest.is_nested(perturbation_scale):
              perturbation_scale = tf.nest.map_structure(
                  lambda x: tf.convert_to_tensor(
                      perturbation_scale,  # pylint: disable=g-long-lambda
                      name="perturbation_scale",
                      dtype=x.dtype,
                  ),
                  unconstrained_params,
              )
          perturbed_unconstrained_parameter_dists = tf.nest.map_structure(
              lambda x, p, s: independent.Independent(  # pylint: disable=g-long-lambda
                  normal.Normal(loc=x, scale=p),
                  reinterpreted_batch_ndims=ps.rank_from_shape(s),
              ),
              unconstrained_params,
              perturbation_scale,
              parameter_prior.event_shape_tensor(),
          )
  
          # For the joint transition, pass the perturbed parameters
          # into the original transition fn (after pushing them into constrained
          # space).
          return joint_distribution_named.JointDistributionNamed(
              ParametersAndState(
                  unconstrained_parameters=_maybe_build_joint_distribution(
                      perturbed_unconstrained_parameter_dists
                  ),
                  state=lambda unconstrained_parameters: (  # pylint: disable=g-long-lambda
                      parameterized_transition_fn(
                          step,
                          state,
                          parameters=parameter_constraining_bijector.forward(
                              unconstrained_parameters
                          ),
                          **kwargs,
                      )
                  ),
              )
          )
  
      return params_and_state_transition_fn
  
imports: import collections
  import contextlib
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.mcmc.particle_filter import infer_trajectories
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 3 for calling a global function with example usage: augment_transition_fn_with_parameters(parameter_prior, parameterized_transition_fn, parameter_constraining_bijector) and returns: params_and_state_transition_fn, joint_distribution_named