code: def get_system_state(
      graph: jraph.GraphsTuple,
  ) -> Tuple[np.ndarray, np.ndarray]:
      return graph.nodes["position"], graph.nodes["momentum"]
  
imports: import functools
  from typing import Tuple, Callable
  from absl import app
  from frozendict import frozendict
  import jax
  import jax.numpy
  import jraph
  import matplotlib.pyplot
  import numpy
label_desc: Use image node_id 6 for calling a global function with example usage: get_system_state(graph) and returns: graph,  graph