class_docstr: JIT (Just-In-Time) module takes a function or a module and compiles it for faster execution.
class_docstr_tok: ["JIT","(","Just-In-Time",")","module","takes","a","function","or","a","module","and","compiles","it","for","faster","execution","."]
docstr: Jit constructor.
  
  Args:
      f: the function or the module to compile.
      vc: the VarCollection of variables used by the function or module. This argument is required for functions.
      static_argnums: tuple of indexes of f's input arguments to treat as static (constants)).
          A new graph is compiled for each different combination of values for such inputs.
docstr_tok: ["Jit","constructor",".","Args",":","f",":","the","function","or","the","module","to","compile",".","vc",":","the","VarCollection","of","variables","used","by","the","function","or","module",".","This","argument","is","required","for","functions",".","static_argnums",":","tuple","of","indexes","of","f","'s","input","arguments","to","treat","as","static","(","constants",")",")",".","A","new","graph","is","compiled","for","each","different","combination","of","values","for","such","inputs","."]
code: def __init__(
      self,
      f: Union[Module, Callable],
      vc: Optional[VarCollection] = None,
      static_argnums: Optional[Tuple[int, ...]] = None,
  ):
      """Jit constructor.
  
      Args:
          f: the function or the module to compile.
          vc: the VarCollection of variables used by the function or module. This argument is required for functions.
          static_argnums: tuple of indexes of f's input arguments to treat as static (constants)).
              A new graph is compiled for each different combination of values for such inputs.
      """
      self.static_argnums = static_argnums
      if not isinstance(f, Module):
          if vc is None:
              raise ValueError(
                  "You must supply the VarCollection used by the function f."
              )
          f = Function(f, vc)
  
      def jit(tensor_list: List[JaxArray], kwargs, *args):
          original_values = self.vc.tensors()
          try:
              self.vc.assign(tensor_list)
              return f(*args, **kwargs), self.vc.tensors()
          finally:
              self.vc.assign(original_values)
  
      self.vc = f.vars() if vc is None else vc
      self._call = jax.jit(
          jit,
          static_argnums=tuple(
              x + 2 for x in sorted(static_argnums or ())
          ),
      )
      self.__wrapped__ = f
  
imports: from collections import namedtuple
  from typing import Optional, List, Union, Callable, Tuple
  import jax
  import jax.numpy
  import numpy
  from objax.typing import JaxArray
  from objax.util import class_name, override_args_kwargs, positional_args_names, repr_function
  from objax.variable import BaseVar, RandomState, VarCollection
label_desc: Use image node_id 1 to create a new Jit object from inherited base classes: Module with example: obj = Jit(f, vc, static_argnums)