class_docstr: A Graph Convolution Network + Pooling model defined with Jraph.
class_docstr_tok: ["A","Graph","Convolution","Network","+","Pooling","model","defined","with","Jraph","."]
docstr: Pooling operation, taken from Jraph.
docstr_tok: ["Pooling","operation",",","taken","from","Jraph","."]
code: def pool(self, graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:
      """Pooling operation, taken from Jraph."""
  
      # Equivalent to jnp.sum(n_node), but JIT-able.
      sum_n_node = graphs.nodes.shape[
          0
      ]  # pytype: disable=attribute-error  # jax-ndarray
      # To aggregate nodes from each graph to global features,
      # we first construct tensors that map the node to the corresponding graph.
      # Example: if you have `n_node=[1,2]`, we construct the tensor [0, 1, 1].
      n_graph = graphs.n_node.shape[0]
      node_graph_indices = jnp.repeat(
          jnp.arange(n_graph),
          graphs.n_node,
          axis=0,
          total_repeat_length=sum_n_node,
      )
      # We use the aggregation function to pool the nodes per graph.
      pooled = self.pooling_fn(
          graphs.nodes, node_graph_indices, n_graph
      )  # pytype: disable=wrong-arg-types  # jax-ndarray
      return graphs._replace(globals=pooled)
  
file_path: flax/examples/ogbg_molpcba/models.py
imports: from typing import Callable, Sequence
  from flax import linen
  import jax.numpy
  import jraph
label_desc: Use image node_id 1 for calling the GraphConvNet obj's underlying member method code with example usage: obj.pool(graphs) and returns: graphs