code: def get_hamiltonian_from_state_fn(
      static_graph: jraph.GraphsTuple,
      hamiltonian_from_graph_fn: Callable[
          [jraph.GraphsTuple], jraph.GraphsTuple
      ],
  ) -> Callable[[np.ndarray, np.ndarray], float]:
      """Returns fn such that fn(position, momentum) -> scalar Hamiltonian.
  
      Args:
          static_graph: `GraphsTuple` containing per-particle static parameters and
             connectivity, such as a full graph of the state can be build by calling
             `set_system_state(static_graph, position, momentum)`.
          hamiltonian_from_graph_fn: callable that given an input `GraphsTuple`
             returns a `GraphsTuple` with a "hamiltonian" field in the globals.
  
      Returns:
          Function that given a state (position, momentum) returns the scalar
          Hamiltonian.
      """
  
      def hamiltonian_from_state_fn(position, momentum):
          # Note we sum along the batch dimension to get the total energy in the batch
          # so get can easily get the gradient.
          graph = set_system_state(static_graph, position, momentum)
          output_graph = hamiltonian_from_graph_fn(graph)
          return output_graph.globals["hamiltonian"].sum()
  
      return hamiltonian_from_state_fn
  
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 8 for calling a global function with example usage: get_hamiltonian_from_state_fn(static_graph, hamiltonian_from_graph_fn) and returns: hamiltonian_from_state_fn, output_graph