code: def _convert_outputs_to_writes(
      nsteps: int, jaxpr: core.Jaxpr, loop_invar_res: Sequence[bool]
  ) -> tuple[core.Jaxpr, list[core.ShapedArray]]:
      assert not jaxpr.constvars, "Jaxpr shouldn't have constvars."
  
      in_avals = [v.aval for v in jaxpr.invars]  # [i, *orig_ref_avals]
  
      @lu.wrap_init
      def eval_jaxpr(i, *refs):
          # We split the refs into the original input refs and the dummy residual
          # refs.
          orig_refs, residual_refs = split_list(
              refs, [len(in_avals) - 1]
          )
          residual_vals = core.eval_jaxpr(jaxpr, (), i, *orig_refs)
          for res_ref, res_val, loop_invar in zip(
              residual_refs, residual_vals, loop_invar_res
          ):
              if loop_invar:
                  res_ref[()] = res_val
              else:
                  res_ref[i] = res_val
          return []
  
      # TODO(mattjj, sharadmv): better handling of tokens, which don't have shape/dtype
      res_ref_avals: list[core.AbstractValue] = [
          AbstractRef(v.aval)
          if loop_invar
          else AbstractRef(  # pytype: disable=attribute-error
              core.ShapedArray(
                  (
                      nsteps,
                      *v.aval.shape,
                  ),  # pytype: disable=attribute-error
                  v.aval.dtype,
              )
          )  # pytype: disable=attribute-error
          for v, loop_invar in zip(jaxpr.outvars, loop_invar_res)
      ]
      jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(
          eval_jaxpr, [*in_avals, *res_ref_avals]
      )
      assert not consts
      return jaxpr, [
          core.ShapedArray(a.shape, a.dtype) for a in res_ref_avals
      ]
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import operator
  from typing import Any, Callable, Generic, TypeVar
  import jax.numpy
  from jax import lax
  from jax.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax.tree_util import tree_flatten, tree_structure, tree_unflatten, treedef_tuple, tree_map, tree_leaves, PyTreeDef
  from jax._src import ad_util
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src.state.types import ReadEffect, AbstractRef, StateEffect
  from jax._src.state import discharge
  from jax._src.state import primitives
  from jax._src.state import utils
  from jax._src.state import types
  from jax._src.typing import Array
  from jax._src.util import partition_list, merge_lists, safe_map, safe_zip, split_list, split_dict, weakref_lru_cache
  from jax._src.lax.control_flow import loops
  from jax._src.lax.control_flow.common import _abstractify, _initial_style_jaxpr
label_desc: Use image node_id 17 for calling a global function with example usage: _convert_outputs_to_writes(nsteps, jaxpr, loop_invar_res) and returns: jaxpr,  unknown, list