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