class_docstr: Traces a JAX function to dot.
class_docstr_tok: ["Traces","a","JAX","function","to","dot","."]
code: def process_call(self, call_primitive, f, tracers, params):
      assert call_primitive.multiple_results
      if call_primitive in (pjit.pjit_p,) and params.get(
          "inline", False
      ):
          f = _interpret_subtrace(f, self.main)
          vals_out = f.call_wrapped(*[t.val for t in tracers])
          return [DotTracer(self, v) for v in vals_out]
  
      graph = Graph.create(
          title=f"{call_primitive} ({name_or_str(f.f)})"
      )
      graph_stack.peek().subgraphs.append(graph)
      with graph_stack(graph):
          f = _interpret_subtrace(f, self.main)
          vals_out = f.call_wrapped(*[t.val for t in tracers])
          return [DotTracer(self, v) for v in vals_out]
  
imports: import collections
  import contextlib
  import functools
  import html
  from typing import Any, Callable, NamedTuple, Optional
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.core
  from jax.experimental import pjit
  from jax.extend import linear_util
label_desc: Use image node_id 5 for calling the DotTrace obj's underlying member method code with example usage: obj.process_call(call_primitive, f, tracers, params) and returns: unknown, unknown