def __init__(self,
               node_model_fn,
               use_received_edges=True,
               use_sent_edges=False,
               use_nodes=True,
               use_globals=True,
               received_edges_reducer=tf.math.unsorted_segment_sum,
               sent_edges_reducer=tf.math.unsorted_segment_sum,
               name="node_block"):
    """Initializes the NodeBlock module.

    Args:
      node_model_fn: A callable that will be called in the variable scope of
        this NodeBlock and should return a Sonnet module (or equivalent
        callable) to be used as the node model. The returned module should take
        a `Tensor` (of concatenated input features for each node) and return a
        `Tensor` (of output features for each node). Typically, this module
        would input and output `Tensor`s of rank 2, but it may also be input or
        output larger ranks. See the `_build` method documentation for more
        details on the acceptable inputs to this module in that case.
      use_received_edges: (bool, default=True) Whether to condition on
        aggregated edges received by each node.
      use_sent_edges: (bool, default=False) Whether to condition on aggregated
        edges sent by each node.
      use_nodes: (bool, default=True) Whether to condition on node attributes.
      use_globals: (bool, default=True) Whether to condition on global
        attributes.
      received_edges_reducer: Reduction to be used when aggregating received
        edges. This should be a callable whose signature matches
        `tf.math.unsorted_segment_sum`.
      sent_edges_reducer: Reduction to be used when aggregating sent edges.
        This should be a callable whose signature matches
        `tf.math.unsorted_segment_sum`.
      name: The module name.

    Raises:
      ValueError: When fields that are required are missing.
    """

    super(NodeBlock, self).__init__(name=name)

    if not (use_nodes or use_sent_edges or use_received_edges or use_globals):
      raise ValueError("At least one of use_received_edges, use_sent_edges, "
                       "use_nodes or use_globals must be True.")

    self._use_received_edges = use_received_edges
    self._use_sent_edges = use_sent_edges
    self._use_nodes = use_nodes
    self._use_globals = use_globals

    with self._enter_variable_scope():
      self._node_model = node_model_fn()
      if self._use_received_edges:
        if received_edges_reducer is None:
          raise ValueError(
              "If `use_received_edges==True`, `received_edges_reducer` "
              "should not be None.")
        self._received_edges_aggregator = ReceivedEdgesToNodesAggregator(
            received_edges_reducer)
      if self._use_sent_edges:
        if sent_edges_reducer is None:
          raise ValueError(
              "If `use_sent_edges==True`, `sent_edges_reducer` "
              "should not be None.")
        self._sent_edges_aggregator = SentEdgesToNodesAggregator(
            sent_edges_reducer)
Initializes the NodeBlock module.

    Args:
      node_model_fn: A callable that will be called in the variable scope of
        this NodeBlock and should return a Sonnet module (or equivalent
        callable) to be used as the node model. The returned module should take
        a `Tensor` (of concatenated input features for each node) and return a
        `Tensor` (of output features for each node). Typically, this module
        would input and output `Tensor`s of rank 2, but it may also be input or
        output larger ranks. See the `_build` method documentation for more
        details on the acceptable inputs to this module in that case.
      use_received_edges: (bool, default=True) Whether to condition on
        aggregated edges received by each node.
      use_sent_edges: (bool, default=False) Whether to condition on aggregated
        edges sent by each node.
      use_nodes: (bool, default=True) Whether to condition on node attributes.
      use_globals: (bool, default=True) Whether to condition on global
        attributes.
      received_edges_reducer: Reduction to be used when aggregating received
        edges. This should be a callable whose signature matches
        `tf.math.unsorted_segment_sum`.
      sent_edges_reducer: Reduction to be used when aggregating sent edges.
        This should be a callable whose signature matches
        `tf.math.unsorted_segment_sum`.
      name: The module name.

    Raises:
      ValueError: When fields that are required are missing.
    
graph_nets/blocks.py