code: def check_jax_usage(enabled: bool = True) -> bool:
      """Ensures JAX APIs (e.g. :func:`jax.vmap`) are used correctly with Haiku.
  
      JAX transforms (like :func:`jax.vmap`) and control flow (e.g.
      :func:`jax.lax.cond`) expect pure functions to be passed in. Some functions
      in Haiku (for example :func:`~haiku.get_parameter`) have side effects and thus
      functions using them are only pure after using :func:`~haiku.transform` (et
      al).
  
      Sometimes it is convenient to use JAX transforms or control flow before
      transforming your function (for example, to :func:`~haiku.vmap` the
      application of a module) but when doing so you need to be careful to use the
      Haiku overloaded version of the underlying JAX function, which carefully makes
      the function(s) you pass in pure functions before calling the underlying JAX
      function.
  
      :func:`check_jax_usage` enables checking raw JAX transforms are used
      appropriately inside Haiku transformed functions. Incorrect usage of JAX
      transforms will result in an error.
  
      Consider the function below, it is not a pure function (a function of its
      inputs with no side effects) because we call into a Haiku API
      (:func:`~haiku.get_parameter`) which during init will create a parameter and
      register it with Haiku.
  
      >>> def f():
      ...   return hk.get_parameter("some_param", [], init=jnp.zeros)
  
      We should not use this with JAX APIs like :func:`jax.vmap` (because it is not
      a pure function). :func:`check_jax_usage` allows you to tell Haiku to make
      incorrect usages of JAX APIs an error:
  
      >>> previous_value = hk.experimental.check_jax_usage(True)
      >>> jax.vmap(f, axis_size=2)()
      Traceback (most recent call last):
        ...
      haiku.JaxUsageError: ...
  
      Using the Haiku wrapped version works correctly:
  
      >>> print(hk.vmap(f, axis_size=2, split_rng=False)())
      [0. 0.]
  
      Args:
        enabled: Boolean indicating whether usage should be checked or not.
  
      Returns:
        Boolean with the previous value for this setting.
      """
      config = get_config()
      previous_value, config.check_jax_usage = (
          config.check_jax_usage,
          enabled,
      )
      return previous_value
  
file_path: dm-haiku/haiku/_src/config.py
imports: import contextlib
  import dataclasses
  import threading
  from typing import Optional
label_desc: Use image node_id 8 for calling a global function with example usage: check_jax_usage(enabled) and returns: previous_value