code: def _get_flat_unconstraining_bijector(jd_model):
      """Create a bijector from a joint distribution that flattens and unconstrains.
  
      The intention is (loosely) to go from a model joint distribution supported on
  
      U_1 x U_2 x ... U_n, with U_j a subset of R^{n_j}
  
      to a model supported on R^N, with N = sum(n_j). (This is "loose" in the sense
      of base measures: some distribution may be supported on an m-dimensional
      subset of R^n, and the default transform for that distribution may then
      have support on R^m. See [1] for details.
  
      Args:
        jd_model: subclass of `tfd.JointDistribution` A JointDistribution for a
          model.
  
      Returns:
        Two `tfb.Bijector`s where the `.forward` method flattens and unconstrains
        points, and the second may be used to initialize a step size.
      """
      # TODO(b/180396233): This bijector is in general point-dependent.
      event_space_bij = (
          jd_model.experimental_default_event_space_bijector()
      )
      flat_bijector = restructure.pack_sequence_as(
          jd_model.event_shape_tensor()
      )
  
      unconstrained_shapes = event_space_bij(
          flat_bijector
      ).inverse_event_shape_tensor(jd_model.event_shape_tensor())
  
      # this reshaping is required as as split can produce a tensor of shape [1]
      # when the distribution event shape is []
      unsplit = joint_map.JointMap(
          tf.nest.map_structure(
              lambda x: reshape.Reshape(
                  event_shape_out=x, event_shape_in=[-1]
              ),
              unconstrained_shapes,
          )
      )
  
      bij = invert.Invert(
          chain.Chain([event_space_bij, flat_bijector, unsplit])
      )
      step_size_bij = invert.Invert(flat_bijector)
  
      return bij, step_size_bij
  
imports: import collections
  import functools
  import warnings
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import restructure
  from tensorflow_probability.python.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import initialization
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.mcmc import preconditioned_nuts
  from tensorflow_probability.python.experimental.mcmc import preconditioning_utils
  from tensorflow_probability.python.experimental.mcmc import sharded
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling a global function with example usage: _get_flat_unconstraining_bijector(jd_model) and returns: bij,  step_size_bij