class_docstr: Represents a function `f` to which `transforms` are to be applied.
  
  Args:
    f: the function to be transformed.
    transforms: a list of `(gen, gen_static_args)` tuples representing
      transformations to apply to `f.` Here `gen` is a generator function and
      `gen_static_args` is a tuple of static arguments for the generator. See
      description at the start of this module for the expected behavior of the
      generator.
    stores: a list of out_store for the auxiliary output of the `transforms`.
    params: extra parameters to pass as keyword arguments to `f`, along with the
      transformed keyword arguments.
class_docstr_tok: ["Represents","a","function","`","f","`","to","which","`","transforms","`","are","to","be","applied",".","Args",":","f",":","the","function","to","be","transformed",".","transforms",":","a","list","of","`","(","gen",",","gen_static_args",")","`","tuples","representing","transformations","to","apply","to","`","f.","`","Here","`","gen","`","is","a","generator","function","and","`","gen_static_args","`","is","a","tuple","of","static","arguments","for","the","generator",".","See","description","at","the","start","of","this","module","for","the","expected","behavior","of","the","generator",".","stores",":","a","list","of","out_store","for","the","auxiliary","output","of","the","`","transforms","`",".","params",":","extra","parameters","to","pass","as","keyword","arguments","to","`","f","`",",","along","with","the","transformed","keyword","arguments","."]
docstr: Calls the underlying function, applying the transforms.
  
  The positional `args` and keyword `kwargs` are passed to the first
  transformation generator.
docstr_tok: ["Calls","the","underlying","function",",","applying","the","transforms",".","The","positional","`","args","`","and","keyword","`","kwargs","`","are","passed","to","the","first","transformation","generator","."]
code: def call_wrapped(self, *args, **kwargs):
      """Calls the underlying function, applying the transforms.
  
      The positional `args` and keyword `kwargs` are passed to the first
      transformation generator.
      """
      stack = []
      for (gen, gen_static_args), out_store in zip(
          self.transforms, self.stores
      ):
          gen = gen(*(gen_static_args + tuple(args)), **kwargs)
          args, kwargs = next(gen)
          stack.append((gen, out_store))
      gen = gen_static_args = out_store = None
  
      try:
          ans = self.f(*args, **dict(self.params, **kwargs))
      except:
          # Some transformations yield from inside context managers, so we have to
          # interrupt them before reraising the exception. Otherwise they will only
          # get garbage-collected at some later time, running their cleanup tasks
          # only after this exception is handled, which can corrupt the global
          # state.
          while stack:
              stack.pop()[0].close()
          raise
  
      args = kwargs = None
      while stack:
          gen, out_store = stack.pop()
          try:
              ans = gen.send(ans)
          except:
              # As above does for the first half of the transformation, exceptions
              # raised in the second half of the transformation also require us to
              # clean up references here.
              while stack:
                  stack.pop()[0].close()
              raise
          if out_store is not None:
              ans, side = ans
              out_store.store(side)
  
      return ans
  
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 5 for calling the WrappedFun obj's underlying member method code with example usage: obj.call_wrapped() and returns: ans