code: def get_higgs_problem(
      min_n_photons: int, max_n_photons: int
  ) -> Problem:
      """Creates fully connected graph containing the detected photons.
  
      Args:
        min_n_photons: minimum number of photons in the detector.
        max_n_photons: maximum number of photons in the detector.
  
      Returns:
        graph, one-hot label whether a higgs was present or not.
      """
      assert min_n_photons >= 2, "Number of photons must be at least 2."
      n_photons = random.randint(min_n_photons, max_n_photons)
      photons = np.stack(
          [get_random_background_photon() for _ in range(n_photons)]
      )
  
      # Add a higgs
      if random.random() > 0.5:
          label = np.eye(2)[0]
          photons[:2] = np.stack(get_random_higgs_photons())
      else:
          label = np.eye(2)[1]
  
      # The graph is fully connected.
      senders = np.repeat(np.arange(n_photons), n_photons)
      receivers = np.tile(np.arange(n_photons), n_photons)
      graph = jraph.GraphsTuple(
          n_node=np.asarray([n_photons]),
          n_edge=np.asarray([len(senders)]),
          nodes=photons,
          edges=None,
          globals=None,
          senders=senders,
          receivers=receivers,
      )
  
      # In order to jit compile our code, we have to pad the nodes and edges of
      # the GraphsTuple to a static shape.
      graph = jraph.pad_with_graphs(
          graph, max_n_photons + 1, max_n_photons * max_n_photons
      )
  
      return Problem(graph=graph, labels=label)
  
file_path: jraph/jraph/examples/higgs_detection.py
imports: import collections
  import logging
  import random
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import jraph
  import numpy
  import optax
  import scipy.stats
label_desc: Use image node_id 5 for calling a global function with example usage: get_higgs_problem(min_n_photons, max_n_photons) and returns: Problem