class_docstr: A `Linker` that JIT-compiles NumPy-based operations using JAX.
code: def jit_compile(self, fn):
      import jax
  
      # I suppose we can consider `Constant`s to be "static" according to
      # JAX.
      static_argnums = [
          n
          for n, i in enumerate(self.fgraph.inputs)
          if isinstance(i, Constant)
      ]
      return jax.jit(fn, static_argnums=static_argnums)
  
file_path: pytensor/pytensor/link/jax/linker.py
imports: import warnings
  from numpy.random import Generator, RandomState
  from pytensor.compile.sharedvalue import SharedVariable, shared
  from pytensor.graph.basic import Constant
  from pytensor.link.basic import JITLinker
label_desc: Use image node_id 2 for calling the JAXLinker obj's underlying member method code with example usage: obj.jit_compile(fn) and returns: jax