code: def trace(state, fn, num_steps, unroll, **_):
      """Implementation of `trace` operator, without the calling convention."""
      # We need the shapes and dtypes of the outputs of `fn`.
      _, untraced_spec, traced_spec = jax.eval_shape(
          fn,
          map_tree(
              lambda s: jax.ShapeDtypeStruct(s.shape, s.dtype), state
          ),
      )
      untraced_init = map_tree(
          lambda spec: jnp.zeros(spec.shape, spec.dtype), untraced_spec
      )
  
      try:
          num_steps = int(num_steps)
          use_scan = True
      except TypeError:
          use_scan = False
          if flatten_tree(traced_spec):
              raise ValueError(
                  "Cannot trace values when `num_steps` is not statically known. Pass "
                  "False to `trace_mask` or return an empty structure (e.g. `()`) as "
                  "the extra output."
              )
          if unroll:
              raise ValueError(
                  "Cannot unroll when `num_steps` is not statically known."
              )
  
      if unroll:
          traced_lists = map_tree(lambda _: [], traced_spec)
          untraced = untraced_init
          for _ in range(num_steps):
              state, untraced, traced_element = fn(state)
              map_tree_up_to(
                  traced_spec,
                  lambda l, e: l.append(e),
                  traced_lists,
                  traced_element,
              )
          # Using asarray instead of stack to handle empty arrays correctly.
          traced = map_tree_up_to(
              traced_spec,
              lambda l, s: jnp.asarray(l, dtype=s.dtype),
              traced_lists,
              traced_spec,
          )
      elif use_scan:
  
          def wrapper(state_untraced, _):
              state, _ = state_untraced
              state, untraced, traced = fn(state)
              return (state, untraced), traced
  
          (state, untraced), traced = lax.scan(
              wrapper,
              (state, untraced_init),
              xs=None,
              length=num_steps,
          )
      else:
          trace_arrays = map_tree(
              lambda spec: jnp.zeros(
                  (num_steps,) + spec.shape, spec.dtype
              ),
              traced_spec,
          )
  
          def wrapper(i, state_untraced_traced):
              state, _, trace_arrays = state_untraced_traced
              state, untraced, traced = fn(state)
              trace_arrays = map_tree(
                  lambda a, e: a.at[i].set(e), trace_arrays, traced
              )
              return (state, untraced, trace_arrays)
  
          state, untraced, traced = lax.fori_loop(
              jnp.asarray(0, num_steps.dtype),
              num_steps,
              wrapper,
              (state, untraced_init, trace_arrays),
          )
      return state, untraced, traced
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  from jax import tree_util
  import jax.numpy
label_desc: Use image node_id 15 for calling a global function with example usage: trace(state, fn, num_steps, unroll) and returns: state,  untraced,  traced, state,  untraced,  traced, state,  untraced,  trace_arrays