def _build(self, graph, edge_model_kwargs=None):
    """Connects the edge block.

    Args:
      graph: A `graphs.GraphsTuple` containing `Tensor`s, whose individual edges
        features (if `use_edges` is `True`), individual nodes features (if
        `use_receiver_nodes` or `use_sender_nodes` is `True`) and per graph
        globals (if `use_globals` is `True`) should be concatenable on the last
        axis.
      edge_model_kwargs: Optional keyword arguments to pass to the `edge_model`.

    Returns:
      An output `graphs.GraphsTuple` with updated edges.

    Raises:
      ValueError: If `graph` does not have non-`None` receivers and senders, or
        if `graph` has `None` fields incompatible with the selected `use_edges`,
        `use_receiver_nodes`, `use_sender_nodes`, or `use_globals` options.
    """
    if edge_model_kwargs is None:
      edge_model_kwargs = {}

    _validate_graph(
        graph, (SENDERS, RECEIVERS, N_EDGE), " when using an EdgeBlock")

    edges_to_collect = []

    if self._use_edges:
      _validate_graph(graph, (EDGES,), "when use_edges == True")
      edges_to_collect.append(graph.edges)

    if self._use_receiver_nodes:
      edges_to_collect.append(broadcast_receiver_nodes_to_edges(graph))

    if self._use_sender_nodes:
      edges_to_collect.append(broadcast_sender_nodes_to_edges(graph))

    if self._use_globals:
      num_edges_hint = _get_static_num_edges(graph)
      edges_to_collect.append(
          broadcast_globals_to_edges(graph, num_edges_hint=num_edges_hint))

    collected_edges = tf.concat(edges_to_collect, axis=-1)
    updated_edges = self._edge_model(collected_edges, **edge_model_kwargs)
    return graph.replace(edges=updated_edges)
Connects the edge block.

    Args:
      graph: A `graphs.GraphsTuple` containing `Tensor`s, whose individual edges
        features (if `use_edges` is `True`), individual nodes features (if
        `use_receiver_nodes` or `use_sender_nodes` is `True`) and per graph
        globals (if `use_globals` is `True`) should be concatenable on the last
        axis.
      edge_model_kwargs: Optional keyword arguments to pass to the `edge_model`.

    Returns:
      An output `graphs.GraphsTuple` with updated edges.

    Raises:
      ValueError: If `graph` does not have non-`None` receivers and senders, or
        if `graph` has `None` fields incompatible with the selected `use_edges`,
        `use_receiver_nodes`, `use_sender_nodes`, or `use_globals` options.
    
graph_nets/blocks.py