Coverage for qml_essentials / tape.py: 98%

1from __future__ import annotations 

2 

3import threading 

4from contextlib import contextmanager 

5from typing import TYPE_CHECKING, Callable, Iterator, List, Optional 

6 

7if TYPE_CHECKING: 

8 from qml_essentials.operations import Operation 

9 

10_local = threading.local() 

11 

12 

13def _tape_stack() -> List[List["Operation"]]: 

14 """Return the per-thread tape stack, creating it on first access. 

15 

16 Returns: 

17 The tape stack for the current thread (a list of tape lists). 

18 """ 

19 if not hasattr(_local, "stack"): 

20 _local.stack = [] 

21 return _local.stack 

22 

23 

24def active_tape() -> Optional[List["Operation"]]: 

25 """Return the innermost active tape, or ``None`` if not recording. 

26 

27 This is called from :meth:`Operation.__init__` to decide whether an 

28 operation should be appended to a tape. 

29 

30 Returns: 

31 The currently active tape list, or ``None``. 

32 """ 

33 stack = _tape_stack() 

34 return stack[-1] if stack else None 

35 

36 

37@contextmanager 

38def recording() -> Iterator[List["Operation"]]: 

39 """Context manager that creates a fresh tape for recording operations. 

40 

41 Operations instantiated inside this block will be appended to the 

42 returned tape list (via :func:`active_tape`). Nesting is supported: 

43 each ``with recording()`` pushes a new tape onto the per-thread stack, 

44 and the previous tape is restored on exit. 

45 

46 Yields: 

47 A new empty list that will be populated with ``Operation`` instances. 

48 """ 

49 stack = _tape_stack() 

50 tape: List["Operation"] = [] 

51 stack.append(tape) 

52 try: 

53 yield tape 

54 finally: 

55 stack.pop() 

56 

57 

58def _pulse_tape_stack() -> List[list]: 

59 """Return the per-thread pulse-event tape stack.""" 

60 if not hasattr(_local, "pulse_stack"): 

61 _local.pulse_stack = [] 

62 return _local.pulse_stack 

63 

64 

65def active_pulse_tape() -> Optional[list]: 

66 """Return the innermost active pulse-event tape, or ``None``. 

67 

68 Called from :class:`~qml_essentials.gates.PulseGates` leaf methods 

69 to record :class:`~qml_essentials.drawing.PulseEvent` objects. 

70 """ 

71 stack = _pulse_tape_stack() 

72 return stack[-1] if stack else None 

73 

74 

75@contextmanager 

76def pulse_recording() -> Iterator[list]: 

77 """Context manager that collects pulse events emitted by PulseGates. 

78 

79 Yields: 

80 A list that will be populated with 

81 :class:`~qml_essentials.drawing.PulseEvent` instances. 

82 """ 

83 stack = _pulse_tape_stack() 

84 tape: list = [] 

85 stack.append(tape) 

86 try: 

87 yield tape 

88 finally: 

89 stack.pop() 

90 

91 

92def shift_and_append(tape_ops: List["Operation"], offset: int) -> None: 

93 """Re-register tape_ops on the active tape with wires shifted by offset. 

94 

95 Each operation is shallow-copied so that the original tape is not 

96 mutated. This is useful for constructing multi-register circuits 

97 where the same sub-circuit must be placed on different qubit 

98 registers. 

99 

100 Args: 

101 tape_ops: List of :class:`Operation` instances (typically captured 

102 via :func:`recording`). 

103 offset: Integer added to every wire index of every operation. 

104 """ 

105 current = active_tape() 

106 if current is None: 

107 return 

108 for o in tape_ops: 

109 shifted = o.__class__.__new__(o.__class__) 

110 shifted.__dict__.update(o.__dict__) 

111 shifted._wires = [w + offset for w in o.wires] 

112 current.append(shifted) 

113 

114 

115def copy_to_tape(fn: Callable, offset: int) -> None: 

116 """Record *fn* into a side tape and replay it shifted onto the active tape. 

117 

118 This is a convenience wrapper around :func:`recording` and 

119 :func:`shift_and_append`. It captures every operation emitted by 

120 *fn* on a temporary tape, then appends shifted copies (wires 

121 incremented by *offset*) to the currently active tape. 

122 

123 Typical usage inside a circuit function:: 

124 

125 def my_circuit(params, inputs, ...): 

126 # first copy on wires 0..n-1 (recorded directly) 

127 model._variational(params, inputs, ...) 

128 # second copy shifted to wires n..2n-1 

129 copy_to_tape(lambda: model._variational(params, inputs, ...), offset=n) 

130 

131 Args: 

132 fn: Zero-argument callable whose body instantiates ``Operation`` 

133 objects (they will be captured on the side tape). 

134 offset: Integer added to every wire index before replaying. 

135 """ 

136 with recording() as side_tape: 

137 fn() 

138 shift_and_append(side_tape, offset)