code: def get_multi_head_attention_example_inputs(
      query_input_dims,
      value_input_dims,
      ragged_key=False,
      ragged_value=False,
      ragged_query=False,
  ):
      """Generates example MultiHeadAttention concrete inputs for testing."""
      # Each batched input is a reshape of a `tf.range()` call.
      batch_size = 2
      # Query input tensor.
      query_size = tf.reduce_prod(query_input_dims)
      query_tsr = tf.range(
          batch_size * query_size, dtype=tf.float32
      ) / tf.cast(query_size, tf.float32)
      query_batch = tf.reshape(
          query_tsr, [batch_size] + query_input_dims
      )
      # Value input tensor.
      value_size = tf.reduce_prod(value_input_dims)
      value_tsr = (
          2.0
          * tf.range(batch_size * value_size, dtype=tf.float32)
          / tf.cast(value_size, tf.float32)
      )
      value_batch = tf.reshape(
          value_tsr, [batch_size] + value_input_dims
      )
      # Key input tensor (optional).
      key_tsr = (
          3.0
          * tf.range(batch_size * value_size, dtype=tf.float32)
          / tf.cast(value_size, tf.float32)
      )
      key_batch = tf.reshape(key_tsr, [batch_size] + value_input_dims)
      # Attention mask input tensor (optional).
      mask_size = tf.reduce_prod(
          query_input_dims[:-1]
      ) * tf.reduce_prod(value_input_dims[:-1])
      mask_input_dims = query_input_dims[:-1] + value_input_dims[:-1]
      mask_tsr = tf.random.uniform([int(batch_size * mask_size)]) <= 0.5
      mask_batch = tf.reshape(mask_tsr, [batch_size] + mask_input_dims)
      # Convert to ragged if needed.
      if ragged_query:
          query_batch = tf.RaggedTensor.from_tensor(query_batch)
      if ragged_value:
          value_batch = tf.RaggedTensor.from_tensor(value_batch)
      if ragged_key:
          key_batch = tf.RaggedTensor.from_tensor(key_batch)
      return query_batch, value_batch, key_batch, mask_batch
  
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 7 for calling a global function with example usage: get_multi_head_attention_example_inputs(query_input_dims, value_input_dims, ragged_key, ragged_value, ragged_query) and returns: query_batch,  value_batch,  key_batch,  mask_batch