code: def _process_eqn(
      eqn: jax.core.JaxprEqn,
      seen: set[str],
      eqns_by_output: Mapping[str, jax.core.JaxprEqn],
      compute_flops: Optional[ComputeFlopsFn],
      scope: _ModuleScope,
      module: Module,
      binder_idx: dict[jax.core.Var, int],
  ) -> Optional[int]:
      """Recursive walks the JaxprEqn to compute the flops it takes."""
      for out_var in eqn.outvars:
          _mark_seen(binder_idx, seen, out_var, scope)
  
      outvars = sorted(
          [_var_to_str(binder_idx, e) for e in eqn.outvars],
          key=_var_sort_key,
      )
      name_stack = str(eqn.source_info.name_stack)
      expression = Expression(
          primitive=eqn.primitive.name,
          invars=" ".join(
              _var_to_str(binder_idx, v) for v in eqn.invars
          )
          if len(eqn.invars) < 10
          else f"{len(eqn.invars)} inputs",
          outvars=" ".join(outvars)
          if len(outvars) < 10
          else f"{outvars[0]}, {len(outvars) - 1} more outputs",
          first_outvar=outvars[0],
          name_stack=name_stack.split("/") if name_stack else [],
      )
  
      if eqn.primitive.name in [
          "named_call",
          "custom_jvp_call_jaxpr",
          "custom_vjp_call_jaxpr",
          "custom_jvp_call",
          "custom_vjp_call",
          "remat_call",
          "scan",
          "while",
          "xla_call",
          "xla_pmap",
          "pjit",
          "remat2",
      ]:
          flops_multiplier = 1
          if eqn.primitive.name in ["named_call", "xla_call"]:
              # Haiku module or named scope.
              name = eqn.params["name"]
              jaxpr = eqn.params["call_jaxpr"]
          elif eqn.primitive.name in ["remat_call", "xla_pmap"]:
              name = eqn.primitive.name + " " + eqn.params["name"]
              jaxpr = eqn.params["call_jaxpr"]
          elif eqn.primitive.name == "remat2":
              name = "remat"
              jaxpr = eqn.params["jaxpr"]
          elif eqn.primitive.name == "while":
              name = "while"
              jaxpr = eqn.params["body_jaxpr"].jaxpr
              # The loop also has a `cond_jaxpr` that we don't display for now.
          elif eqn.primitive.name == "scan":
              name = f'scan length={eqn.params["length"]} unroll={eqn.params["unroll"]}'
              jaxpr = eqn.params["jaxpr"].jaxpr
              flops_multiplier = (
                  eqn.params["length"] // eqn.params["unroll"]
              )
          elif eqn.primitive.name == "pjit":
              name = eqn.params["name"]
              jaxpr = eqn.params["jaxpr"].jaxpr
          elif eqn.primitive.name in [
              "custom_jvp_call",
              "custom_vjp_call",
          ]:
              name = eqn.primitive.name.replace("_call", "")
              jaxpr = eqn.params["call_jaxpr"]
          elif eqn.primitive.name in [
              "custom_jvp_call_jaxpr",
              "custom_vjp_call_jaxpr",
          ]:
              # Custom gradient.
              name = eqn.primitive.name.replace("_call_jaxpr", "")
              jaxpr = eqn.params["fun_jaxpr"].jaxpr
          else:
              raise ValueError(f"unmatched eqn {eqn.primitive.name}")
  
          expression.submodule = Module(name=name)
          flops = _process_jaxpr(
              jaxpr,
              compute_flops,
              scope.join(eqn),
              seen,
              expression.submodule,
          )
          if compute_flops is not None:
              flops *= flops_multiplier
  
          expression.submodule.flops = flops
          expression.flops = flops
          if compute_flops is None or flops > 0:
              module.expressions.append(expression)
      else:
          details = []
          if eqn.invars:
              details.append(
                  "in "
                  + ", ".join(_format_shape(v) for v in eqn.invars)
              )
          if eqn.outvars:
              details.append(
                  "out "
                  + ", ".join(_format_shape(v) for v in eqn.outvars)
              )
  
          expression.details = ", ".join(details)
          expression.params.update(
              {k: str(v) for k, v in eqn.params.items()}
          )
  
          flops = (
              None
              if compute_flops is None
              else compute_flops(eqn, expression)
          )
          expression.flops = flops
          module.expressions.append(expression)
  
      for var in eqn.invars:
          if isinstance(var, jax.core.Literal):
              continue
  
          key = _var_to_str(binder_idx, var)
          if key == "*":
              continue
  
          if not _mark_seen(binder_idx, seen, var, scope):
              continue
  
          if key not in eqns_by_output:
              logging.warning("missing var %s = %s", key, type(var))
              continue
          f = _process_eqn(
              eqns_by_output[key],
              seen,
              eqns_by_output,
              compute_flops,
              scope,
              module,
              binder_idx,
          )
          if compute_flops is not None:
              flops += f
  
      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 7 for calling a global function with example usage: _process_eqn(eqn, seen, eqns_by_output, compute_flops, scope, module, binder_idx) and returns: flops