code: def _process_jaxpr(
      jaxpr: jax.core.Jaxpr,
      compute_flops: Optional[ComputeFlopsFn],
      scope: _ModuleScope,
      seen: set[str],
      module: Module,
  ) -> Optional[int]:
      """Computes the flops used for a JAX expression, tracking module scope."""
      if isinstance(jaxpr, jax.core.ClosedJaxpr):
          return _process_jaxpr(
              jaxpr.jaxpr, compute_flops, scope, seen, module
          )
  
      # Label variables by the order in which they're introduced.
      lam_binders = itertools.chain(jaxpr.constvars, jaxpr.invars)
      let_binders = itertools.chain.from_iterable(
          e.outvars for e in jaxpr.eqns
      )
      all_binders = itertools.chain(lam_binders, let_binders)
      binder_idx = dict(zip(all_binders, itertools.count()))
  
      # Build a map that for each output contains the equation to compute it.
      eqns_by_output = {}
      for eqn in jaxpr.eqns:
          for var in eqn.outvars:
              eqns_by_output[_var_to_str(binder_idx, var)] = eqn
  
      # Seed the set of variables that have been seen with the inputs and constants.
      for var in jaxpr.invars + jaxpr.constvars:
          _mark_seen(binder_idx, seen, var, scope)
  
      # Recursively walk the computation graph.
      flops = None if compute_flops is None else 0
      for var in jaxpr.outvars:
          if isinstance(var, jax.core.Var) and _mark_seen(
              binder_idx, seen, var, scope
          ):
              f = _process_eqn(
                  eqns_by_output[_var_to_str(binder_idx, var)],
                  seen,
                  eqns_by_output,
                  compute_flops,
                  scope,
                  module,
                  binder_idx,
              )
              if compute_flops is not None:
                  flops += f
  
      module.expressions.sort(
          key=lambda e: _var_sort_key(e.first_outvar)
      )
      module.flops = flops
      return flops
  
imports: from collections.abc import Mapping, Sequence
  import dataclasses
  import itertools
  import logging
  import os
  import sys
  from typing import Any, Callable, NamedTuple, Optional
  from haiku._src import summarise
  import jax
  import jax.core
label_desc: Use image node_id 8 for calling a global function with example usage: _process_jaxpr(jaxpr, compute_flops, scope, seen, module) and returns: flops, _process_jaxpr