code: def without_apply_rng(f: TransformedT) -> TransformedT:
      """Removes the rng argument from the apply function.
  
      This is a convenience wrapper that makes the ``rng`` argument to
      ``f.apply`` default to ``None``. This is useful when ``f`` doesn't actually
      use random numbers as part of its computation, such that the ``rng`` argument
      wouldn't be used. Note that if ``f`` `does` use random numbers, this will
      cause an error to be thrown complaining that ``f`` needs a non-None PRNGKey.
  
      Args:
        f: A transformed function.
  
      Returns:
        The same transformed function, with a modified ``apply``.
      """
  
      def check_rng_kwarg(kwargs):
          if "rng" in kwargs:
              raise TypeError(
                  "Haiku transform adds three arguments (params, state, rng) to apply. "
                  "If the functions you are transforming use the same names you must "
                  "pass them positionally (e.g. `f.apply(.., my_rng)` and not by "
                  "name (e.g. `f.apply(.., rng=my_rng)`)"
              )
  
      if isinstance(f, hk.TransformedWithState):
  
          def apply_fn(params, state, *args, **kwargs):
              check_rng_kwarg(kwargs)
              return f.apply(params, state, None, *args, **kwargs)
  
          apply_fn.__signature__ = (
              _transform.sig_replace_leading_parameters(
                  inspect.signature(f.apply),
                  3,
                  [
                      inspect.Parameter(
                          "params",
                          inspect.Parameter.POSITIONAL_OR_KEYWORD,
                          annotation=Optional[hk.Params],
                      ),
                      inspect.Parameter(
                          "state",
                          inspect.Parameter.POSITIONAL_OR_KEYWORD,
                          annotation=Optional[hk.State],
                      ),
                  ],
              )
          )
          f_new = hk.TransformedWithState(init=f.init, apply=apply_fn)
          _transform.tie_in_original_fn(f, f_new.init, f_new.apply)
  
      elif isinstance(f, hk.Transformed):
  
          def apply_fn(params, *args, **kwargs):
              check_rng_kwarg(kwargs)
              return f.apply(params, None, *args, **kwargs)
  
          apply_fn.__signature__ = (
              _transform.sig_replace_leading_parameters(
                  inspect.signature(f.apply),
                  2,
                  [
                      inspect.Parameter(
                          "params",
                          inspect.Parameter.POSITIONAL_OR_KEYWORD,
                          annotation=Optional[hk.Params],
                      ),
                  ],
              )
          )
          f_new = hk.Transformed(init=f.init, apply=apply_fn)
          _transform.tie_in_original_fn(f, f_new.init, f_new.apply)
  
      elif isinstance(f, MultiTransformedWithState):
  
          def make_new_apply_fn(
              apply_fn, params, state, *args, **kwargs
          ):
              check_rng_kwarg(kwargs)
              return apply_fn(params, state, None, *args, **kwargs)
  
          apply_fn = jax.tree_util.tree_map(
              lambda fn: functools.partial(make_new_apply_fn, fn),
              f.apply,
          )
          f_new = MultiTransformedWithState(init=f.init, apply=apply_fn)
  
      elif isinstance(f, MultiTransformed):
  
          def make_new_apply_fn(apply_fn, params, *args, **kwargs):
              check_rng_kwarg(kwargs)
              return apply_fn(params, None, *args, **kwargs)
  
          apply_fn = jax.tree_util.tree_map(
              lambda fn: functools.partial(make_new_apply_fn, fn),
              f.apply,
          )
          f_new = MultiTransformed(init=f.init, apply=apply_fn)
  
      else:
          raise ValueError(
              "Must be called with the result of `hk.transformed`, "
              "`hk.multi_transform`, `hk.transformed_with_state` or "
              "`hk.multi_transform_with_state`, "
              f"actually called with {type(f)}"
          )
  
      return f_new
  
file_path: dm-haiku/haiku/_src/multi_transform.py
imports: import dataclasses
  import functools
  import inspect
  from typing import Any, Callable, NamedTuple, Optional, TypeVar
  from haiku._src import analytics
  from haiku._src import transform
  from haiku._src import typing
  import jax
label_desc: Use image node_id 4 for calling a global function with example usage: without_apply_rng(f) and returns: f_new, f, f, apply_fn, apply_fn