class_docstr: A Graph Convolution Network + Pooling model defined with Jraph.
class_docstr_tok: ["A","Graph","Convolution","Network","+","Pooling","model","defined","with","Jraph","."]
code: def __call__(self, graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:
      # We will first linearly project the original node features as 'embeddings'.
      embedder = jraph.GraphMapFeatures(
          embed_node_fn=nn.Dense(self.latent_size)
      )
      processed_graphs = embedder(graphs)
  
      # Now, we will apply the GCN once for each message-passing round.
      for _ in range(self.message_passing_steps):
          mlp_feature_sizes = [self.latent_size] * self.num_mlp_layers
          update_node_fn = jraph.concatenated_args(
              MLP(
                  mlp_feature_sizes,
                  dropout_rate=self.dropout_rate,
                  deterministic=self.deterministic,
              )
          )
          graph_conv = jraph.GraphConvolution(
              update_node_fn=update_node_fn, add_self_edges=True
          )
  
          if self.skip_connections:
              processed_graphs = add_graphs_tuples(
                  graph_conv(processed_graphs), processed_graphs
              )
          else:
              processed_graphs = graph_conv(processed_graphs)
  
          if self.layer_norm:
              processed_graphs = processed_graphs._replace(
                  nodes=nn.LayerNorm()(processed_graphs.nodes),
              )
  
      # We apply the pooling operation to get a 'global' embedding.
      processed_graphs = self.pool(processed_graphs)
  
      # Now, we decode this to get the required output logits.
      decoder = jraph.GraphMapFeatures(
          embed_global_fn=nn.Dense(self.output_globals_size)
      )
      processed_graphs = decoder(processed_graphs)
  
      return processed_graphs
  
file_path: jax/jax/experimental/jax2tf/tests/flax_models/gnn.py
imports: from collections.abc import Sequence
  from typing import Callable
  from flax import linen
  import jax
  import jax.numpy
  import jraph
label_desc: Use image node_id 2 for calling the GraphConvNet obj's underlying member method code with example usage: obj.__call__(graphs) and returns: processed_graphs