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)
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.
    
graph_nets/modules.py