code: def _loop_invariant_outputs(jaxpr: core.Jaxpr) -> list[bool]:
      # Get effects for each of the jaxpr inputs and remove the loop index.
      ref_effects = state_types.get_ref_state_effects(
          [v.aval for v in jaxpr.invars], jaxpr.effects
      )[1:]
      # We first assume that *read-only `Ref`s* are loop-invariant. We can safely do
      # this because the only way something can be loop-varying is if we write to it
      # at some point. It's *possible* that read-write `Ref`s are loop-invariant but
      # we conservatively assume they aren't.
      loop_invar_refs = [
          _is_read_only(effs) if effs else True for effs in ref_effects
      ]
      loop_var_refs = map(operator.not_, loop_invar_refs)
  
      # We'd like to detect if the outputs of the jaxpr are loop-invariant. An
      # output is loop-invariant if it is downstream of only loop-invariant values
      # (seeded by the read-only `Ref`s). If at any point, a loop-varying value
      # interacts with a loop-invariant value, we produce a loop-varying value. We
      # can use `partial_eval` to perform this analysis by treating loop-varying
      # values as "unknown" and loop-invariant values as "known", since when a known
      # and unknown value interact, they produce an unknown value.
      loop_var_inputs = [True, *loop_var_refs]
      (
          _,
          _,
          loop_var_outputs,
          _,
          _,
      ) = _partial_eval_jaxpr_custom(
          jaxpr, loop_var_inputs, _save_everything
      )
      return map(operator.not_, loop_var_outputs)
  
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 14 for calling a global function with example usage: _loop_invariant_outputs(jaxpr) and returns: map