class_docstr: Output type of :func:`pyro.ops.jit.trace`.
  
  Wrapper around the output of :func:`torch.jit.trace`
  that handles parameter plumbing.
  
  The actual PyTorch compilation artifact is stored in :attr:`compiled`.
  Call diagnostic methods on this attribute.
class_docstr_tok: ["Output","type","of",":","func",":","`","pyro.ops.jit.trace","`",".","Wrapper","around","the","output","of",":","func",":","`","torch.jit.trace","`","that","handles","parameter","plumbing",".","The","actual","PyTorch","compilation","artifact","is","stored","in",":","attr",":","`","compiled","`",".","Call","diagnostic","methods","on","this","attribute","."]
code: def __call__(self, *args, **kwargs):
      key = _hashable_args_kwargs(args, kwargs)
  
      # if first time
      if key not in self.compiled:
          # param capture
          with poutine.block():
              with poutine.trace(
                  param_only=True
              ) as first_param_capture:
                  self.fn(*args, **kwargs)
  
          self._param_names = list(
              set(first_param_capture.trace.nodes.keys())
          )
          unconstrained_params = tuple(
              pyro.param(name).unconstrained()
              for name in self._param_names
          )
          params_and_args = unconstrained_params + args
          weakself = weakref.ref(self)
  
          def compiled(*params_and_args):
              self = weakself()
              unconstrained_params = params_and_args[
                  : len(self._param_names)
              ]
              args = params_and_args[len(self._param_names) :]
              constrained_params = {}
              for name, unconstrained_param in zip(
                  self._param_names, unconstrained_params
              ):
                  constrained_param = pyro.param(
                      name
                  )  # assume param has been initialized
                  assert (
                      constrained_param.unconstrained()
                      is unconstrained_param
                  )
                  constrained_params[name] = constrained_param
              return poutine.replay(self.fn, params=constrained_params)(
                  *args, **kwargs
              )
  
          if self.ignore_warnings:
              compiled = ignore_jit_warnings()(compiled)
          with pyro.validation_enabled(False):
              time_compilation = self.jit_options.pop(
                  "time_compilation", False
              )
              with optional(timed(), time_compilation) as t:
                  self.compiled[key] = torch.jit.trace(
                      compiled, params_and_args, **self.jit_options
                  )
              if time_compilation:
                  self.compile_time = t.elapsed
      else:
          unconstrained_params = [
              pyro.param(name).unconstrained()
              for name in self._param_names
          ]
          params_and_args = unconstrained_params + list(args)
  
      with poutine.block(hide=self._param_names):
          with poutine.trace(param_only=True) as param_capture:
              ret = self.compiled[key](*params_and_args)
  
      for name in param_capture.trace.nodes.keys():
          if name not in self._param_names:
              raise NotImplementedError(
                  "pyro.ops.jit.trace assumes all params are created on "
                  "first invocation, but found new param: {}".format(
                      name
                  )
              )
  
      return ret
  
file_path: pyro/pyro/ops/jit.py
imports: import argparse
  import warnings
  import weakref
  import torch
  import pyro
  import pyro.poutine
  from pyro.util import ignore_jit_warnings, optional, timed
label_desc: Use image node_id 2 for calling the CompiledFunction obj's underlying member method code with example usage: obj.__call__() and returns: ret, poutine