code: def transformation_with_aux(
      gen, fun: lu.WrappedFun, *gen_static_args
  ) -> tuple[lu.WrappedFun, Any]:
      out_store = StoreEqual()
      out_thunk = lambda: out_store.val
      return fun.wrap(gen, gen_static_args, out_store), out_thunk
  
imports: from __future__ import annotations
  import functools
  from typing import Any, Callable
  from jax._src import ad_util
  from jax._src import api_util
  from jax._src import core
  from jax._src import custom_api_util
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import traceback_util
  from jax._src import util
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import xla
  from jax._src.tree_util import tree_flatten, tree_leaves, tree_map, tree_structure, treedef_tuple, tree_unflatten
label_desc: Use image node_id 1 for calling a global function with example usage: transformation_with_aux(gen, fun) and returns: fun,  out_thunk