code: def maybe_jit(f, num_args):
      static_argnums = thin(range(num_args), 0.5)
  
      def fun(*args):
          partial_args = list(args)
          for i in static_argnums:
              partial_args[i] = None
  
          @jit
          def jitted_fun(*partial_args):
              full_args = list(partial_args)
              for i in static_argnums:
                  full_args[i] = args[i]
              return f(*full_args)
  
          return jitted_fun(*partial_args)
  
      return fun
  
imports: from collections import namedtuple
  from functools import partial
  import numpy.random
  from absl.testing import absltest
  from absl.testing import parameterized
  import itertools
  import jax.numpy
  from jax import jit, jvp, vjp
  import jax._src.test_util
  from jax import config
  from jax._src.util import unzip2, safe_zip, safe_map
label_desc: Use image node_id 4 for calling a global function with example usage: maybe_jit(f, num_args) and returns: fun, jitted_fun, f