code: def test_verify_consistent_outputs(
      self,
      use_key,
      use_query_as_kwarg,
      use_value_as_kwarg,
      use_attention_mask,
      return_scores,
      ragged_key,
      ragged_value,
      ragged_query,
  ):
      num_heads = 2
      key_dim, value_dim = (2, 3)
      query_input_dims = [2, 3]
      value_input_dims = [3, 4]
      (
          query_batch,
          value_batch,
          key_batch,
          mask_batch,
      ) = get_multi_head_attention_example_inputs(
          query_input_dims,
          value_input_dims,
          ragged_key=ragged_key,
          ragged_value=ragged_value,
          ragged_query=ragged_query,
      )
      layer_instance = tf.keras.layers.MultiHeadAttention(
          num_heads, key_dim, value_dim
      )
  
      # Set up test inputs, branched on input order.
      input_kwargs = {
          "key": key_batch if use_key else None,
          "attention_mask": mask_batch if use_attention_mask else None,
          "return_attention_scores": return_scores,
      }
      input_args = tuple()
      if use_value_as_kwarg and use_query_as_kwarg:
          input_kwargs["query"] = query_batch
          input_kwargs["value"] = value_batch
      elif use_value_as_kwarg:
          input_kwargs["value"] = value_batch
          input_args = (query_batch,)
      elif use_query_as_kwarg:
          # Invalid test case; cannot pass query kwarg after value.
          return
      else:
          input_args = (query_batch, key_batch)
  
      dummy_tape = tf.GradientTape()
      with dummy_tape:
          (
              _,
              computed_outputs,
              _,
          ) = multi_head_attention.multi_head_attention_layer_computation(
              layer_instance=layer_instance,
              input_args=input_args,
              input_kwargs=input_kwargs,
              tape=dummy_tape,
          )
      true_outputs = layer_instance(*input_args, **input_kwargs)
      self.assertAllClose(computed_outputs, true_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 1 for calling the CheckOutputs obj's underlying member method code with example usage: obj.test_verify_consistent_outputs(use_key, use_query_as_kwarg, use_value_as_kwarg, use_attention_mask, return_scores, ragged_key, ragged_value, ragged_query) without return types