code: def make_two_layer_attention_model(
      layer_generator, input_dims, output_dims
  ):
      """Creates a 2-layer MultiHeadAttention model."""
      (
          inputs,
          input_args,
          input_kwargs,
      ) = get_multi_head_attention_model_inputs(input_dims)
      layer1 = layer_generator(input_dims, output_dims)
      temp1 = layer1(*input_args, **input_kwargs)
      temp2 = tf.keras.layers.Dense(1)(temp1)
      outputs = common_test_utils.reshape_and_sum(temp2)
      return tf.keras.Model(inputs=inputs, outputs=outputs)
  
imports: from absl.testing import parameterized
  import tensorflow
  from tensorflow_privacy.privacy.fast_gradient_clipping import common_test_utils
  from tensorflow_privacy.privacy.fast_gradient_clipping import layer_registry
  from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import dense
  from tensorflow_privacy.privacy.fast_gradient_clipping.registry_functions import multi_head_attention
label_desc: Use image node_id 3 for calling a global function with example usage: make_two_layer_attention_model(layer_generator, input_dims, output_dims) and returns: tf