docstr: Test a combination of transforms.
docstr_tok: ["Test","a","combination","of","transforms","."]
code: def test_wrapped_fun_transforms(self):
      """Test a combination of transforms."""
  
      def f(*args, **kwargs):
          """The function to be transformed.
          Scales the positional arguments by a factor.
          Takes only one keyword argument, the factor to scale by."""
          factor = kwargs.pop("factor", 2)  # For PY2
          assert not kwargs
          return tuple(a * factor for a in args)
  
      @lu.transformation_with_aux
      def kw_to_positional(factor, *args, **kwargs):
          """A transformation with auxiliary output.
          Turns all keyword parameters into positional ones.
  
          On entry, append the values of the keyword arguments to the positional
          arguments. On exit, take a list of results and recreate a dictionary
          from the tail of the results. The auxiliary output is the list of
          keyword keys.
          """
          kwargs_keys = kwargs.keys()
          new_args = tuple(kwargs[k] for k in kwargs_keys)
          new_kwargs = dict(factor=factor)
          results = (
              yield args + new_args,
              new_kwargs,
          )  # Yield transformed (args, kwargs)
          # Assume results correspond 1:1 to the args + new_args
          assert len(results) == len(args) + len(new_args)
          aux_output = len(new_args)
          yield (
              results[0 : len(args)],
              dict(zip(kwargs_keys, results[len(args) :])),
          ), aux_output
  
      wf = lu.wrap_init(f)  # Wraps `f` as a `WrappedFun`.
      wf, out_thunk = kw_to_positional(wf, 2)
      # Call the transformed function.
      scaled_positional, scaled_kwargs = wf.call_wrapped(
          1, 2, three=3, four=4
      )
      self.assertEqual((2, 4), scaled_positional)
      self.assertEqual(dict(three=6, four=8), scaled_kwargs)
      self.assertEqual(2, out_thunk())
  
file_path: jax/tests/util_test.py
imports: import operator
  from absl.testing import absltest
  from jax._src import linear_util
  from jax._src import test_util
  from jax._src import util
  from jax import config
  from jax._src.util import weakref_lru_cache
label_desc: Use image node_id 1 for calling the UtilTest obj's underlying member method code with example usage: obj.test_wrapped_fun_transforms() and returns: tuple