code: def transformation_with_aux(
      gen, fun: WrappedFun, *gen_static_args, use_eq_store=False
  ) -> tuple[WrappedFun, Any]:
      """Adds one more transformation with auxiliary output to a WrappedFun."""
      out_store = Store() if not use_eq_store else EqualStore()
      out_thunk = lambda: out_store.val
      return fun.wrap(gen, gen_static_args, out_store), out_thunk
  
imports: from __future__ import annotations
  from functools import partial
  import operator
  from typing import Any, Callable, NamedTuple
  import weakref
  from jax._src import config
  from jax._src import core
  from jax._src import traceback_util
  from jax._src.tree_util import tree_map
  from jax._src.util import curry
label_desc: Use image node_id 2 for calling a global function with example usage: transformation_with_aux(gen, fun) and returns: fun,  out_thunk