A class representing a single invocation of an operator. It binds the operator instance and the call arguments. It also tracks the order of invocations of stateful operators, which is important for lazy evaluation of stateful operators or operators with side-effects.
| 32 | |
| 33 | |
| 34 | class Invocation: |
| 35 | """ |
| 36 | A class representing a single invocation of an operator. |
| 37 | |
| 38 | It binds the operator instance and the call arguments. |
| 39 | It also tracks the order of invocations of stateful operators, which is important for |
| 40 | lazy evaluation of stateful operators or operators with side-effects. |
| 41 | """ |
| 42 | |
| 43 | def __init__( |
| 44 | self, |
| 45 | operator_instance: "Operator", |
| 46 | call_id: Optional[int], |
| 47 | inputs: list[Any] | None = None, |
| 48 | args: dict[str, Any] | None = None, |
| 49 | is_batch: bool = False, |
| 50 | batch_size: Optional[int] = None, |
| 51 | previous_invocation: Optional["Invocation"] = None, |
| 52 | caller_frame: types.FrameType | None = None, |
| 53 | ): |
| 54 | """ |
| 55 | Parameters |
| 56 | ---------- |
| 57 | operator_instance : OperatorInstance |
| 58 | The operator instance that is being invoked. |
| 59 | call_id : int |
| 60 | The call ID of the invocation - necessary to avoid folding of multiple invocations |
| 61 | of the same stateful operator. |
| 62 | inputs : list |
| 63 | The inputs to the operator. |
| 64 | args : dict name->argument |
| 65 | The argument inputs of the operator. Scalar arguments are part of the operator instance. |
| 66 | is_batch : bool |
| 67 | Whether this is a batch invocation. |
| 68 | NOTE: A batch of 1 and a single tensor are equivalent from the operator's perspective |
| 69 | (operators always work with batches) but differes from the user's perspective. |
| 70 | batch_size : int |
| 71 | The batch sizes. This is useful chiefly for operators without inputs or ones that alter |
| 72 | the batch size. |
| 73 | previous_invocation : Invocation |
| 74 | The previous invocation of the same operator. Used by stateful operators. |
| 75 | caller_frame : FrameType |
| 76 | The resolved user call-site frame. Used to capture call stacks for error reporting. |
| 77 | """ |
| 78 | self._operator = operator_instance |
| 79 | self._call_id = call_id |
| 80 | self._inputs = inputs or [] |
| 81 | self._args = args or {} |
| 82 | self._is_batch = is_batch |
| 83 | self._results: tuple[Any] | None = None |
| 84 | self._batch_size = batch_size |
| 85 | self._num_outputs: int | None = None |
| 86 | self._output_devices: list[Device] | None = None |
| 87 | self._previous_invocation = previous_invocation |
| 88 | self._eval_context = _EvalContext.current()._snapshot() |
| 89 | self._eval_mode: _EvalMode | None = None |
| 90 | self._future: Optional[_Future] = None |
| 91 | self._run_lock = threading.Lock() |