code: def _make_vector_event_space_bijector(model):
      """Creates a vector bijector that constrains like the structured model."""
  
      # TODO(siege): Make this work with multi-part default_event_bijector.
      def _make_reshaped_bijector(b, s):
          return tfb.Chain(
              [
                  tfb.Reshape(
                      event_shape_in=s,
                      event_shape_out=[ps.reduce_prod(s)],
                  ),
                  b,
                  tfb.Reshape(
                      event_shape_in=[
                          ps.reduce_prod(b.inverse_event_shape(s))
                      ],
                      event_shape_out=b.inverse_event_shape(s),
                  ),
              ]
          )
  
      reshaped_bijector = tf.nest.map_structure(
          _make_reshaped_bijector,
          model.default_event_space_bijector,
          model.event_shape,
      )
  
      return tfb.Blockwise(
          bijectors=tf.nest.flatten(reshaped_bijector),
          block_sizes=tf.nest.flatten(
              tf.nest.map_structure(
                  lambda b, s: ps.reduce_prod(
                      b.inverse_event_shape(s)
                  ),  # pylint: disable=g-long-lambda
                  model.default_event_space_bijector,
                  model.event_shape,
              )
          ),
      )
  
file_path: probability/spinoffs/inference_gym/inference_gym/targets/vector_model.py
imports: import collections
  import functools
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import prefer_static
  from inference_gym.targets import model
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling a global function with example usage: _make_vector_event_space_bijector(model) and returns: tfb, tfb