code: def GraphMapFeatures(
      embed_edge_fn: Optional[EmbedEdgeFn] = None,
      embed_node_fn: Optional[EmbedNodeFn] = None,
      embed_global_fn: Optional[EmbedGlobalFn] = None,
  ):
      """Returns function which embeds the components of a graph independently.
  
      Args:
        embed_edge_fn: function used to embed the edges.
        embed_node_fn: function used to embed the nodes.
        embed_global_fn: function used to embed the globals.
      """
      identity = lambda x: x
      embed_edges_fn = embed_edge_fn if embed_edge_fn else identity
      embed_nodes_fn = embed_node_fn if embed_node_fn else identity
      embed_global_fn = embed_global_fn if embed_global_fn else identity
  
      def Embed(graphs_tuple):
          return graphs_tuple._replace(
              nodes=embed_nodes_fn(graphs_tuple.nodes),
              edges=embed_edges_fn(graphs_tuple.edges),
              globals=embed_global_fn(graphs_tuple.globals),
          )
  
      return Embed
  
file_path: jraph/jraph/_src/models.py
imports: import functools
  from typing import Any, Callable, Iterable, Mapping, Optional, Union
  import jax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import utils
label_desc: Use image node_id 3 for calling a global function with example usage: GraphMapFeatures(embed_edge_fn, embed_node_fn, embed_global_fn) and returns: Embed, graphs_tuple