class_docstr: Traces a JAX function to dot.
class_docstr_tok: ["Traces","a","JAX","function","to","dot","."]
code: def process_primitive(self, primitive, tracers, params):
      val_out = primitive.bind(*[t.val for t in tracers], **params)
      if primitive is pjit.pjit_p:
          f = jax.core.jaxpr_as_fun(params["jaxpr"])
          f.__name__ = params["name"]
          fun = lu.wrap_init(f)
          return self.process_call(primitive, fun, tracers, params)
  
      inputs = [t.val for t in tracers]
      outputs = list(jax.tree_util.tree_leaves(val_out))
  
      graph = graph_stack.peek()
      node = Node(id=outputs[0], title=str(primitive), outputs=outputs)
      graph.nodes.append(node)
      graph.edges.extend([(i, outputs[0]) for i in inputs])
  
      return jax.tree_util.tree_map(
          lambda v: DotTracer(self, v), val_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 4 for calling the DotTrace obj's underlying member method code with example usage: obj.process_primitive(primitive, tracers, params) and returns: jax, self