code: def model_forward_pass(
      input_model: tf.keras.Model,
      inputs: type_aliases.PackedTensors,
      generator_fn: type_aliases.GeneratorFunction = None,
  ) -> tuple[type_aliases.PackedTensors, Sequence[Any]]:
      """Does a forward pass of a model and returns useful intermediates.
  
      NOTE: the graph traversal algorithm is an adaptation of the logic in the
        _run_internal_graph() method in the functional.Functional class. Hence,
        forward_norm_pass should only be invoked if the generated model
        instance is an instance of the functional.Functional class.
  
      Args:
        input_model: A `tf.keras.Model` to compute the quantities for.
        inputs: Arbitrary input to be fed into the input layer of the model. It is
          expected that `input_model(inputs)` returns a valid output.
        generator_fn: A function with signature `(tf.keras.layers.Layer, Any, Any)
          -> (tf.Tensor, Any)`, where we require `generator_fn(layer_instance, args,
          kwargs)[0] == layer_instance(*args, **kwargs)`. If `None`, then
          `layer_fn(layer_instance, args, kwargs)[1] == None`.
  
      Returns:
        A `tuple` `(outputs, generator_outputs_list)`. `outputs` is the
        `PackedTensor` that is generated as a result of a forward pass.
        `generator_outputs_list` is a `list` whose i-th entry is the output of
        `generator_fn(lyr, args, kwargs)[1]` where `lyr` is the i-th
        layer when the compute graph of `input_model` is traversed in BFS order.
      """
      # TODO: Avoid or remove the references to protected methods of `input_model`.  # pylint: disable=g-bad-todo
  
      # Default generator.
      generator_outputs_list = []
      if generator_fn is None:
  
          def generator_fn(layer_instance, args, kwargs):
              return layer_instance(*args, **kwargs), None
  
      # Prepare the inputs and BFS variables.
      flattened_inputs = input_model._flatten_to_reference_inputs(
          inputs
      )  # pylint: disable=protected-access
      tensor_dict = {}
      tensor_usage_count = (
          input_model._tensor_usage_count
      )  # pylint: disable=protected-access
      for x, y in zip(input_model.inputs, flattened_inputs):
          y = input_model._conform_to_reference_input(
              y, ref_input=x
          )  # pylint: disable=protected-access
          x_id = str(id(x))
          tensor_dict[x_id] = [y] * tensor_usage_count[x_id]
      nodes_by_depth = (
          input_model._nodes_by_depth
      )  # pylint: disable=protected-access
      depth_keys = list(nodes_by_depth.keys())
      depth_keys.sort(reverse=True)
  
      # Perform BFS feedforward computations.
      for depth in depth_keys:
          for node in nodes_by_depth[depth]:
              if node.is_input:
                  continue  # inputs already exist
              if any(
                  t_id not in tensor_dict
                  for t_id in node.flat_input_ids
              ):
                  continue  # node is not computable; try skipping
              args, kwargs = node.map_arguments(tensor_dict)
              if has_internal_compute_graph(node.layer):
                  # If this node has an internal computational graph, we can recurse.
                  (
                      node_layer_outputs,
                      node_generator_outputs,
                  ) = model_forward_pass(node.layer, args, generator_fn)
                  generator_outputs_list.extend(node_generator_outputs)
              else:
                  # Otherwise, we parse the node directly.
                  (
                      node_layer_outputs,
                      layer_generator_outputs,
                  ) = generator_fn(node.layer, args, kwargs)
                  generator_outputs_list.append(layer_generator_outputs)
  
              # Update the current dictionary of inputs for the next node.
              for x_id, y in zip(
                  node.flat_output_ids,
                  tf.nest.flatten(node_layer_outputs),
              ):
                  tensor_dict[x_id] = [y] * tensor_usage_count[x_id]
  
      # Gather outputs (in case there are multiple) and return.
      output_tensors = []
      for x in input_model.outputs:
          x_id = str(id(x))
          output_tensors.append(tensor_dict[x_id].pop())
      model_outputs = tf.nest.pack_sequence_as(
          input_model._nested_outputs,  # pylint: disable=protected-access
          output_tensors,
      )
  
      return model_outputs, generator_outputs_list
  
file_path: privacy/tensorflow_privacy/privacy/fast_gradient_clipping/gradient_clipping_utils.py
imports: from collections.abc import Sequence, Set
  from typing import Any, Literal, Optional
  from absl import logging
  import tensorflow
  from tensorflow_privacy.privacy.fast_gradient_clipping import layer_registry
  from tensorflow_privacy.privacy.fast_gradient_clipping import type_aliases
label_desc: Use image node_id 2 for calling a global function with example usage: model_forward_pass(input_model, inputs, generator_fn) and returns: model_outputs,  generator_outputs_list, layer_instance,  None