code: def network_definition(graph: jraph.GraphsTuple) -> jraph.ArrayTree:
      """`InteractionNetwork` with an LSTM in the edge update."""
  
      # LSTM that will keep a memory of the inputs to the edge model.
      edge_fn_lstm = hk.LSTM(hidden_size=HIDDEN_SIZE)
  
      # MLPs used in the edge and the node model. Note that in this instance
      # the output size matches the input size so the same model can be run
      # iteratively multiple times. In a real model, this would usually be achieved
      # by first using an encoder in the input data into a common `EMBEDDING_SIZE`.
      edge_fn_mlp = hk.nets.MLP([HIDDEN_SIZE, EMBEDDING_SIZE])
      node_fn_mlp = hk.nets.MLP([HIDDEN_SIZE, EMBEDDING_SIZE])
  
      # Initialize the edge features to contain both the input edge embedding
      # and initial LSTM state. Note for the nodes we only have an embedding since
      # in this example nodes do not use a `node_fn_lstm`, but for analogy, we
      # still put it in a `StatefulField`.
      graph = graph._replace(
          edges=StatefulField(
              embedding=graph.edges,
              state=edge_fn_lstm.initial_state(graph.edges.shape[0]),
          ),
          nodes=StatefulField(embedding=graph.nodes, state=None),
      )
  
      def update_edge_fn(edges, sender_nodes, receiver_nodes):
          # We will run an LSTM memory on the inputs first, and then
          # process the output of the LSTM with an MLP.
          edge_inputs = jnp.concatenate(
              [
                  edges.embedding,
                  sender_nodes.embedding,
                  receiver_nodes.embedding,
              ],
              axis=-1,
          )
          lstm_output, updated_state = edge_fn_lstm(
              edge_inputs, edges.state
          )
          updated_edges = StatefulField(
              embedding=edge_fn_mlp(lstm_output),
              state=updated_state,
          )
          return updated_edges
  
      def update_node_fn(nodes, received_edges):
          # Note `received_edges.state` will also contain the aggregated state for
          # all received edges, which we may choose to use in the node update.
          node_inputs = jnp.concatenate(
              [nodes.embedding, received_edges.embedding], axis=-1
          )
          updated_nodes = StatefulField(
              embedding=node_fn_mlp(node_inputs), state=None
          )
          return updated_nodes
  
      recurrent_graph_network = jraph.InteractionNetwork(
          update_edge_fn=update_edge_fn, update_node_fn=update_node_fn
      )
  
      # Apply the model recurrently for 10 message passing steps.
      # If instead we intended to use the LSTM to process a sequence of features
      # for each node/edge, here we would select the corresponding inputs from the
      # sequence along the sequence axis of the nodes/edges features to build the
      # correct input graph for each step of the iteration.
      num_message_passing_steps = 10
      for _ in range(num_message_passing_steps):
          graph = recurrent_graph_network(graph)
  
      return graph
  
file_path: jraph/jraph/examples/lstm.py
imports: import collections
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import jax.tree_util
  import jraph
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: network_definition(graph) and returns: graph, updated_edges, updated_nodes