code: def lorenz_transition_fn(_, loc, innovation_scale, step_size):
      x, y, z = tf.unstack(loc, axis=-1)
      dx = 10 * (y - x)
      dy = x * (28 - z) - y
      dz = x * y - 8 / 3 * z
      delta = tf.stack([dx, dy, dz], axis=-1)
      return tfd.Independent(
          tfd.Normal(
              loc + step_size * delta,
              tf.sqrt(step_size) * innovation_scale[..., tf.newaxis],
          ),
          reinterpreted_batch_ndims=1,
      )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import tensor_util
  from inference_gym.internal import data
  from inference_gym.targets import bayesian_model
  from inference_gym.targets import model
  from inference_gym.targets.ground_truth import convection_lorenz_bridge
  from inference_gym.targets.ground_truth import convection_lorenz_bridge_unknown_scales
label_desc: Use image node_id 1 for calling a global function with example usage: lorenz_transition_fn(_, loc, innovation_scale, step_size) and returns: tfd