class_docstr: CommNet module.
  
  Implementation for the model originally described in
  https://arxiv.org/abs/1605.07736 (S. Sukhbaatar, A. Szlam, R. Fergus), in the
  version presented in https://arxiv.org/abs/1706.06122 (Y. Hoshen).
  
  This module internally creates edge features based on the features from the
  nodes sending to that edge, and independently learns an embedding for each
  node. It then uses these edges and nodes features to compute updated node
  features.
  
  This module does not use the global nor the edges features of the input, but
  uses its receivers and senders information. The output graph has the same
  value in edge and global fields as the input graph. The edge and global
  features fields may have a `None` value in the input `gn_graphs.GraphsTuple`.
docstr: Initializes the CommNet module.
  
  Args:
    edge_model_fn: A callable to be passed to EdgeBlock. The callable must
      return a Sonnet module (or equivalent; see EdgeBlock for details).
    node_encoder_model_fn: A callable to be passed to the NodeBlock
      responsible for the first encoding of the nodes. The callable must
      return a Sonnet module (or equivalent; see NodeBlock for details). The
      shape of this module's output should match the shape of the module built
      by `edge_model_fn`, but for the first and last dimension.
    node_model_fn: A callable to be passed to NodeBlock. The callable must
      return a Sonnet module (or equivalent; see NodeBlock for details).
    reducer: Reduction to be used when aggregating the edges in the nodes.
      This should be a callable whose signature matches
      tf.math.unsorted_segment_sum.
    name: The module name.
code: def __init__(
      self,
      edge_model_fn,
      node_encoder_model_fn,
      node_model_fn,
      reducer=tf.math.unsorted_segment_sum,
      name="comm_net",
  ):
      """Initializes the CommNet module.
  
      Args:
        edge_model_fn: A callable to be passed to EdgeBlock. The callable must
          return a Sonnet module (or equivalent; see EdgeBlock for details).
        node_encoder_model_fn: A callable to be passed to the NodeBlock
          responsible for the first encoding of the nodes. The callable must
          return a Sonnet module (or equivalent; see NodeBlock for details). The
          shape of this module's output should match the shape of the module built
          by `edge_model_fn`, but for the first and last dimension.
        node_model_fn: A callable to be passed to NodeBlock. The callable must
          return a Sonnet module (or equivalent; see NodeBlock for details).
        reducer: Reduction to be used when aggregating the edges in the nodes.
          This should be a callable whose signature matches
          tf.math.unsorted_segment_sum.
        name: The module name.
      """
      super(CommNet, self).__init__(name=name)
  
      with self._enter_variable_scope():
          # Computes $\Psi_{com}(x_j)$ in Eq. (2) of 1706.06122
          self._edge_block = blocks.EdgeBlock(
              edge_model_fn=edge_model_fn,
              use_edges=False,
              use_receiver_nodes=False,
              use_sender_nodes=True,
              use_globals=False,
          )
          # Computes $\Phi(x_i)$ in Eq. (2) of 1706.06122
          self._node_encoder_block = blocks.NodeBlock(
              node_model_fn=node_encoder_model_fn,
              use_received_edges=False,
              use_sent_edges=False,
              use_nodes=True,
              use_globals=False,
              received_edges_reducer=reducer,
              name="node_encoder_block",
          )
          # Computes $\Theta(..)$ in Eq.(2) of 1706.06122
          self._node_block = blocks.NodeBlock(
              node_model_fn=node_model_fn,
              use_received_edges=True,
              use_sent_edges=False,
              use_nodes=True,
              use_globals=False,
              received_edges_reducer=reducer,
          )
  
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 1 to create a new CommNet object from inherited base classes: _base with example: obj = CommNet(edge_model_fn, node_encoder_model_fn, node_model_fn, reducer, name)