code: def transform_and_run(
      f: Optional[Fn] = None,
      seed: Optional[int] = 42,
      run_apply: bool = True,
      jax_transform: Optional[Callable[[Fn], Fn]] = None,
      *,
      map_rng: Optional[Callable[[Key], Key]] = None,
  ) -> T:
      r"""Transforms the given function and runs init then (optionally) apply.
  
      Equivalent to:
  
      >>> def f(x):
      ...   return x
      >>> x = jnp.ones([])
      >>> rng = jax.random.PRNGKey(42)
      >>> f = hk.transform_with_state(f)
      >>> params, state = f.init(rng, x)
      >>> out = f.apply(params, state, rng, x)
  
      This function makes it very convenient to unit test Haiku:
  
      >>> class MyTest(unittest.TestCase):
      ...   @hk.testing.transform_and_run
      ...   def test_linear_output(self):
      ...     mod = hk.Linear(1)
      ...     out = mod(jnp.ones([1, 1]))
      ...     self.assertEqual(out.ndim, 2)
  
      It can also be combined with ``chex`` to test all pure/jit/pmap versions of a
      function:
  
      >>> class MyTest(unittest.TestCase):
      ...   @chex.all_variants
      ...   def test_linear_output(self):
      ...     @hk.testing.transform_and_run(jax_transform=self.variant)
      ...     def f(inputs):
      ...       mod = hk.Linear(1)
      ...       return mod(inputs)
      ...     out = f(jnp.ones([1, 1]))
      ...     self.assertEqual(out.ndim, 2)
  
      And can also be useful in an interactive environment like ipython, Jupyter or
      Google Colaboratory:
  
      >>> f = lambda x: hk.Bias()(x)
      >>> print(hk.testing.transform_and_run(f)(jnp.ones([1, 1])))
      [[1.]]
  
      See :func:`transform` for more details.
  
      To use this with `pmap` (without ``chex``) you need to additionally pass in a
      function to map the init/apply rng keys. For example, if you want every
      instance of your pmap to have the same key:
  
      >>> def same_key_on_all_devices(key):
      ...   return jnp.broadcast_to(key, (jax.local_device_count(), *key.shape))
  
      >>> @hk.testing.transform_and_run(jax_transform=jax.pmap,
      ...                               map_rng=same_key_on_all_devices)
      ... def test_something():
      ...   ...
  
      Or you can use a different key:
  
      >>> def different_key_on_all_devices(key):
      ...   return jax.random.split(key, jax.local_device_count())
  
      >>> @hk.testing.transform_and_run(jax_transform=jax.pmap,
      ...                               map_rng=different_key_on_all_devices)
      ... def test_something_else():
      ...   ...
  
      Args:
        f: A function method to transform.
        seed: A seed to pass to init and apply.
        run_apply: Whether to run apply as well as init. Defaults to true.
        jax_transform: An optional jax transform to apply on the init and apply
          functions.
        map_rng: If set to a non-None value broadcast the init/apply rngs
          broadcast_rng-ways.
  
      Returns:
        A function that :func:`~haiku.transform`\ s ``f`` and runs ``init`` and
        optionally ``apply``.
      """
      if f is None:
          return functools.partial(
              transform_and_run,
              seed=seed,
              run_apply=run_apply,
              jax_transform=jax_transform,
              map_rng=map_rng,
          )
  
      @functools.wraps(f)
      def wrapper(*a, **k):
          """Runs init and apply of f."""
          if seed is not None:
              init_rng, apply_rng = (
                  jax.random.PRNGKey(seed),
                  jax.random.PRNGKey(seed + 1),
              )
              if map_rng is not None:
                  init_rng, apply_rng = map(
                      map_rng, (init_rng, apply_rng)
                  )
          else:
              init_rng, apply_rng = None, None
          init, apply = transform.transform_with_state(
              lambda: f(*a, **k)
          )
          if jax_transform:
              init, apply = map(jax_transform, (init, apply))
          params, state = init(init_rng)
          if run_apply:
              out, state = apply(params, state, apply_rng)
              return out
  
      return wrapper
  
file_path: dm-haiku/haiku/_src/test_utils.py
imports: from collections.abc import Generator, Sequence
  import functools
  import inspect
  import itertools
  import os
  import types
  from typing import Any, Callable, Optional, TypeVar
  from absl.testing import parameterized
  from haiku._src import config
  from haiku._src import transform
  import jax
label_desc: Use image node_id 1 for calling a global function with example usage: transform_and_run(f, seed, run_apply, jax_transform) and returns: wrapper, functools, out