code: def set_system_state(
      static_graph: jraph.GraphsTuple,
      position: np.ndarray,
      momentum: np.ndarray,
  ) -> jraph.GraphsTuple:
      """Sets the non-static parameters of the graph (momentum, position)."""
      nodes = static_graph.nodes.copy(
          position=position, momentum=momentum
      )
      return static_graph._replace(nodes=nodes)
  
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 5 for calling a global function with example usage: set_system_state(static_graph, position, momentum) and returns: static_graph