code: def when(condition):
      def _wrapped(f):
          if isinstance(condition, bool):
              if condition:
                  f()
          else:
              lax.cond(condition, f, lambda: None)
  
      return _wrapped
  
file_path: jax/jax/_src/pallas/utils.py
imports: import math
  import numpy
  from jax import lax
  from jax._src import core
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: when(condition) and returns: _wrapped