class_docstr: Multi-head self-attention module.
  
  The module is based on the following three papers:
   * A simple neural network module for relational reasoning (RNs):
       https://arxiv.org/abs/1706.01427
   * Non-local Neural Networks: https://arxiv.org/abs/1711.07971.
   * Attention Is All You Need (AIAYN): https://arxiv.org/abs/1706.03762.
  
  The input to the modules consists of a graph containing values for each node
  and connectivity between them, a tensor containing keys for each node
  and a tensor containing queries for each node.
  
  The self-attention step consist of updating the node values, with each new
  node value computed in a two step process:
  - Computing the attention weights between each node and all of its senders
   nodes, by calculating sum(sender_key*receiver_query) and using the softmax
   operation on all attention weights for each node.
  - For each receiver node, compute the new node value as the weighted average
   of the values of the sender nodes, according to the attention weights.
  - Nodes with no received edges, get an updated value of 0.
  
  Values, keys and queries contain a "head" axis to compute independent
  self-attention for each of the heads.
docstr: Connects the multi-head self-attention module.
  
  The self-attention is only computed according to the connectivity of the
  input graphs, with receiver nodes attending to sender nodes.
  
  Args:
    node_values: Tensor containing the values associated to each of the nodes.
      The expected shape is [total_num_nodes, num_heads, key_size].
    node_keys: Tensor containing the key associated to each of the nodes. The
      expected shape is [total_num_nodes, num_heads, key_size].
    node_queries: Tensor containing the query associated to each of the nodes.
      The expected shape is [total_num_nodes, num_heads, query_size]. The
      query size must be equal to the key size.
    attention_graph: Graph containing connectivity information between nodes
      via the senders and receivers fields. Node A will only attempt to attend
      to Node B if `attention_graph` contains an edge sent by Node A and
      received by Node B.
  
  Returns:
    An output `graphs.GraphsTuple` with updated nodes containing the
    aggregated attended value for each of the nodes with shape
    [total_num_nodes, num_heads, value_size].
  
  Raises:
    ValueError: if the input graph does not have edges.
code: def _build(
      self, node_values, node_keys, node_queries, attention_graph
  ):
      """Connects the multi-head self-attention module.
  
      The self-attention is only computed according to the connectivity of the
      input graphs, with receiver nodes attending to sender nodes.
  
      Args:
        node_values: Tensor containing the values associated to each of the nodes.
          The expected shape is [total_num_nodes, num_heads, key_size].
        node_keys: Tensor containing the key associated to each of the nodes. The
          expected shape is [total_num_nodes, num_heads, key_size].
        node_queries: Tensor containing the query associated to each of the nodes.
          The expected shape is [total_num_nodes, num_heads, query_size]. The
          query size must be equal to the key size.
        attention_graph: Graph containing connectivity information between nodes
          via the senders and receivers fields. Node A will only attempt to attend
          to Node B if `attention_graph` contains an edge sent by Node A and
          received by Node B.
  
      Returns:
        An output `graphs.GraphsTuple` with updated nodes containing the
        aggregated attended value for each of the nodes with shape
        [total_num_nodes, num_heads, value_size].
  
      Raises:
        ValueError: if the input graph does not have edges.
      """
  
      # Sender nodes put their keys and values in the edges.
      # [total_num_edges, num_heads, query_size]
      sender_keys = blocks.broadcast_sender_nodes_to_edges(
          attention_graph.replace(nodes=node_keys)
      )
      # [total_num_edges, num_heads, value_size]
      sender_values = blocks.broadcast_sender_nodes_to_edges(
          attention_graph.replace(nodes=node_values)
      )
  
      # Receiver nodes put their queries in the edges.
      # [total_num_edges, num_heads, key_size]
      receiver_queries = blocks.broadcast_receiver_nodes_to_edges(
          attention_graph.replace(nodes=node_queries)
      )
  
      # Attention weight for each edge.
      # [total_num_edges, num_heads]
      attention_weights_logits = tf.reduce_sum(
          sender_keys * receiver_queries, axis=-1
      )
      normalized_attention_weights = _received_edges_normalizer(
          attention_graph.replace(edges=attention_weights_logits),
          normalizer=self._normalizer,
      )
  
      # Attending to sender values according to the weights.
      # [total_num_edges, num_heads, embedding_size]
      attented_edges = (
          sender_values * normalized_attention_weights[..., None]
      )
  
      # Summing all of the attended values from each node.
      # [total_num_nodes, num_heads, embedding_size]
      received_edges_aggregator = blocks.ReceivedEdgesToNodesAggregator(
          reducer=tf.math.unsorted_segment_sum
      )
      aggregated_attended_values = received_edges_aggregator(
          attention_graph.replace(edges=attented_edges)
      )
  
      return attention_graph.replace(nodes=aggregated_attended_values)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from graph_nets import _base
  from graph_nets import blocks
  import tensorflow
label_desc: Use image node_id 2 for calling the SelfAttention obj's underlying member method code with example usage: obj._build(node_values, node_keys, node_queries, attention_graph) and returns: attention_graph