code: def network_definition(graph: jraph.GraphsTuple) -> jraph.ArrayTree:
      """Defines a graph neural network.
  
      Args:
        graph: Graphstuple the network processes.
  
      Returns:
        globals.
      """
  
      @jax.vmap
      @jraph.concatenated_args
      def update_edge_fn(features):
          return hk.nets.MLP([30, 30, 30])(features)
  
      # The correct solution for the edge update function is the invariant mass
      # of the photon pair.
      # The simple MLP we use here seems to fail to find the correct solution.
      # You can ensure that the example works in principle by replacing the
      # update_edge_fn below with the following analytical solution.
      @jax.vmap
      def unused_update_edge_fn_solution(s, r):
          """Calculates invariant mass of photon pair and compares to Higgs mass."""
          t = (s + r) ** 2
          return jnp.array(
              jnp.abs(t[0] - t[1] - t[2] - t[3] - 125.18**2) < 1,
              dtype=jnp.float32,
          )[None]
  
      gn = jraph.RelationNetwork(
          update_edge_fn=update_edge_fn,
          update_global_fn=hk.nets.MLP([2]),
          aggregate_edges_for_globals_fn=jraph.segment_sum,
      )
      graph = gn(graph)
  
      return graph.globals
  
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 6 for calling a global function with example usage: network_definition(graph) and returns: graph, hk, jnp