code: def check_not_jax_transformed(f):
      # TODO(tomhennigan): Consider `CompiledFunction = type(jax.jit(lambda: 0))`.
      if isinstance(f, COMPILED_FN_TYPES):
          raise ValueError(
              "A common error with Haiku is to pass an already jit "
              "(or pmap) decorated function into hk.transform (e.g. "
              "`hk.transform(jax.jit(f)))`. You should instead jit/pmap "
              "the init or apply function you get back from Haiku (e.g. "
              "`jax.jit(hk.transform(f).apply)`).\n\n"
              "This is because the function you pass into hk.transform "
              "is not a pure function (because you don't explicitly "
              "pass in/out params/rng). jit and pmap require you to "
              "pass in a pure function (such as the init or apply "
              "functions Haiku gives you back from hk.transform)."
          )
  
imports: from collections.abc import Mapping
  import inspect
  from typing import Any, Callable, NamedTuple, Optional, TypeVar, Union
  from haiku._src import analytics
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import typing
  import jax
label_desc: Use image node_id 8 for calling a global function with example usage: check_not_jax_transformed(f) without return types