code: def _flatten_nested_jacobian(jacobian, state_shape):
      """Flattens a nested Jacobian into a matrix.
  
      The flattening and concatenation follows the interpretation of the structure
      as being a leading 'axis', meaning that if the input has 'shape':
      [input_structure, A, B], and the output has 'shape':
      [output_structure, C, D], the input Jacobian should have the 'shape':
      [input_structure, output_structure, A, B, C, D]. As with the regular axes, the
      encoding is input major.
  
      Args:
        jacobian: A nested Jacobian.
        state_shape: A nested collection of state shapes.
  
      Returns:
        jacobian_mat: The Jacobian matrix.
  
      #### Examples
  
      Non-structured state:
  
      ```python
      input = tf.zeros([1, 2])
      output = tf.zeros([3])
      jacobian = tf.zeros([1, 2, 3])
      ```
  
      Structured state:
  
      ```python
      input = {'x': tf.zeros([1, 2])}
      output = {'y': tf.zeros([3])}
      jacobian = {'x': {'y': tf.zeros([1, 2, 3])}}
      ```
  
      A more complicated structure:
  
      ```python
      input = [tf.zeros([1, 2]), tf.zeros([])]
      output = {'y': tf.zeros([3])}
      jacobian = [{'y': tf.zeros([1, 2, 3])}, {'y': tf.zeros([3]}]
      ```
  
      """
  
      def _flatten_row(jacobian_row, state_shape_part):
          state_size = ps.reduce_prod(state_shape_part)
          jacobian_row_mats = tf.nest.map_structure(
              lambda j: tf.reshape(
                  j, ps.stack([state_size, -1], axis=0)
              ),
              jacobian_row,
          )
          return tf.concat(tf.nest.flatten(jacobian_row_mats), axis=-1)
  
      flat_rows = nest.map_structure_up_to(
          state_shape, _flatten_row, jacobian, state_shape
      )
      return tf.concat(tf.nest.flatten(flat_rows), axis=0)
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math import gradient
  from tensorflow.python.util import nest
label_desc: Use image node_id 6 for calling a global function with example usage: _flatten_nested_jacobian(jacobian, state_shape) and returns: tf, tf