class_docstr: A complete Graph Network model defined with Jraph.
class_docstr_tok: ["A","complete","Graph","Network","model","defined","with","Jraph","."]
code: def __call__(self, graphs: jraph.GraphsTuple) -> jraph.GraphsTuple:
      # We will first linearly project the original features as 'embeddings'.
      embedder = jraph.GraphMapFeatures(
          embed_node_fn=nn.Dense(self.latent_size),
          embed_edge_fn=nn.Dense(self.latent_size),
          embed_global_fn=nn.Dense(self.latent_size),
      )
      processed_graphs = embedder(graphs)
  
      # Now, we will apply a Graph Network once for each message-passing round.
      mlp_feature_sizes = [self.latent_size] * self.num_mlp_layers
      for _ in range(self.message_passing_steps):
          if self.use_edge_model:
              update_edge_fn = jraph.concatenated_args(
                  MLP(
                      mlp_feature_sizes,
                      dropout_rate=self.dropout_rate,
                      deterministic=self.deterministic,
                  )
              )
          else:
              update_edge_fn = None
  
          update_node_fn = jraph.concatenated_args(
              MLP(
                  mlp_feature_sizes,
                  dropout_rate=self.dropout_rate,
                  deterministic=self.deterministic,
              )
          )
          update_global_fn = jraph.concatenated_args(
              MLP(
                  mlp_feature_sizes,
                  dropout_rate=self.dropout_rate,
                  deterministic=self.deterministic,
              )
          )
  
          graph_net = jraph.GraphNetwork(
              update_node_fn=update_node_fn,
              update_edge_fn=update_edge_fn,
              update_global_fn=update_global_fn,
          )
  
          if self.skip_connections:
              processed_graphs = add_graphs_tuples(
                  graph_net(processed_graphs), processed_graphs
              )
          else:
              processed_graphs = graph_net(processed_graphs)
  
          if self.layer_norm:
              processed_graphs = processed_graphs._replace(
                  nodes=nn.LayerNorm()(processed_graphs.nodes),
                  edges=nn.LayerNorm()(processed_graphs.edges),
                  globals=nn.LayerNorm()(processed_graphs.globals),
              )
  
      # Since our graph-level predictions will be at globals, we will
      # decode 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 1 for calling the GraphNet obj's underlying member method code with example usage: obj.__call__(graphs) and returns: processed_graphs