Coverage for qml_essentials / model.py: 92%
705 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-09-03 21:15 +0000
1from typing import Any, Dict, Optional, Tuple, Callable, Union, List
3import warnings
4import jax.numpy as jnp
5import numpy as np
6import jax
7from jax import random
9import jaqsi as js
10from jaqsi import operations as op
11from jaqsi import gateset
12from jaqsi.tape import recording
13from jaqsi.noise import KrausChannel
14from jaqsi import noise
15from jaqsi.gates import Gates, PulseInformation as pinfo
16from jaqsi import make_hashable, safe_random_split
18from qml_essentials.ansaetze import Ansaetze, Circuit, Encoding
20import logging
22log = logging.getLogger(__name__)
24# model mode -> (ansatz/state-prep gates at pulse level, encoding gates at pulse level)
25GATE_MODES = {
26 "unitary": (False, False),
27 "ansatz_pulse": (True, False),
28 "enc_pulse": (False, True),
29 "all_pulse": (True, True),
30}
32# the modes that run the respective gate group at pulse level. Both only serve
33# the deprecated explicit gate_mode argument and can go with it.
34_ANSATZ_PULSE_MODES = ("ansatz_pulse", "all_pulse")
35_ENC_PULSE_MODES = ("enc_pulse", "all_pulse")
38class Model:
39 """
40 A quantum circuit model.
41 """
43 def __init__(
44 self,
45 n_qubits: int,
46 n_layers: int,
47 circuit_type: Union[str, type[Circuit]] = "No_Ansatz",
48 data_reupload: Union[
49 bool, List[List[bool]], List[List[List[bool]]], np.ndarray
50 ] = True,
51 state_preparation: Union[
52 str, Callable, List[Union[str, Callable]], None
53 ] = None,
54 encoding: Union[Encoding, str, Callable, List[Union[str, Callable]]] = Gates.RX,
55 trainable_frequencies: bool = False,
56 initialization: str = "random",
57 initialization_domain: List[float] = [0, 2 * jnp.pi],
58 output_qubit: Union[List[int], int, None] = None,
59 observables: Union[
60 int, List[Union[int, List[int]]], List[op.Operation], None
61 ] = None,
62 shots: Optional[int] = None,
63 random_seed: int = 1000,
64 repeat_batch_axis: List[bool] = [True, True, True, True],
65 pulse_shape: str = "gaussian",
66 ) -> None:
67 """
68 Initialize the quantum circuit model.
69 Parameters will have the shape [impl_n_layers, parameters_per_layer]
70 where impl_n_layers is the number of layers provided and added by one
71 depending if data_reupload is True and parameters_per_layer is given by
72 the chosen ansatz.
74 The model is initialized with the following parameters as defaults:
75 - noise_params: None
76 - execution_type: "expval"
77 - shots: None
79 Args:
80 n_qubits (int): The number of qubits in the circuit.
81 n_layers (int): The number of layers in the circuit.
82 circuit_type (str, Circuit): The type of quantum circuit to use.
83 If None, defaults to "no_ansatz".
84 data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
85 np.ndarray], optional):
86 Whether to reupload data to the quantum device on each
87 layer and qubit. Detailed re-uploading instructions can be given
88 as a list/array of 0/False and 1/True with shape (n_qubits,
89 n_layers) to specify where to upload the data. Defaults to True
90 for applying data re-uploading to the full circuit.
91 encoding (Union[str, Callable, List[str], List[Callable]], optional):
92 The unitary to use for encoding the input data. Can be a string
93 (e.g. "RX") or a callable (e.g. gateset.RX). Defaults to gateset.RX.
94 If input is multidimensional it is assumed to be a list of
95 unitaries or a list of strings.
96 trainable_frequencies (bool, optional):
97 Sets trainable encoding parameters for trainable frequencies.
98 Defaults to False.
99 initialization (str, optional): The strategy to initialize the parameters.
100 Can be "random", "zeros", "zero-controlled", "pi", or "pi-controlled".
101 Defaults to "random".
102 output_qubit (List[int], int, optional): Deprecated alias for
103 ``observables``. Forwards to ``observables`` and will be removed
104 in a future release. Defaults to None.
105 observables (int, List[int], List[List[int]], List[op.Operation],
106 optional): Measurement specification. A qubit index, a list of
107 indices, or a list of qubit groups (for $Z$-parity) selects the
108 measured subsystem with the default PauliZ readout.
109 Alternatively, a list of
110 :class:`~jaqsi.operations.Operation` observables makes
111 ``execution_type="expval"`` return one expectation value per
112 observable. When None all qubits are measured. Defaults to None.
113 shots (Optional[int], optional): The number of shots to use for
114 the quantum device. Defaults to None.
115 random_seed (int, optional): seed for the random number generator
116 in initialization is "random" and for random noise parameters.
117 Defaults to 1000.
118 repeat_batch_axis (List[bool], optional): Each boolean in the array
119 determines over which axes to parallelise computation. The axes
120 correspond to [inputs, params, pulse_params, enc_pulse_params].
121 Defaults to [True, True, True, True], meaning that batching is
122 enabled over all axes. A 3-element list (legacy) is accepted and
123 extended with a trailing True for the enc_pulse_params axis.
124 pulse_shape (str, optional): Pulse envelope shape for pulse-level
125 simulation. One of ``PulseEnvelope.available()``.
126 Defaults to ``"gaussian"``.
128 Returns:
129 None
130 """
131 # Initialize default parameters needed for circuit evaluation
132 self.n_qubits: int = n_qubits
133 if output_qubit is not None:
134 if observables is not None:
135 raise ValueError("Pass either output_qubit or observables, not both.")
136 warnings.warn(
137 "output_qubit is deprecated, use observables instead.",
138 DeprecationWarning,
139 stacklevel=2,
140 )
141 observables = output_qubit
142 self.observables = observables
143 self.n_layers: int = n_layers
144 self.noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None
145 self.shots = shots
146 self.trainable_frequencies: bool = trainable_frequencies
147 self.execution_type: str = "expval"
148 # backward compatibility
149 # TODO: consider making this more generic in future
150 # (for someone wanting to control this without bothering with pulse stuff)
151 if len(repeat_batch_axis) == 3:
152 log.warning("Batch axis should have length 4")
153 repeat_batch_axis = list(repeat_batch_axis) + [True]
154 self.repeat_batch_axis: List[bool] = repeat_batch_axis
156 # --- Pulse envelope ---
157 pinfo.set_envelope(pulse_shape)
159 # --- State Preparation ---
160 try:
161 self._sp = Gates.parse_gates(state_preparation, Gates)
162 except ValueError as e:
163 raise ValueError(f"Error parsing encodings: {e}")
165 # prepare corresponding pulse parameters (always optimized pulses)
166 self.sp_pulse_params = []
167 for sp in self._sp:
168 sp_name = sp.__name__ if hasattr(sp, "__name__") else str(sp)
170 if pinfo.gate_by_name(sp_name) is not None:
171 self.sp_pulse_params.append(pinfo.gate_by_name(sp_name).params)
172 else:
173 # gate has no pulse parametrization
174 self.sp_pulse_params.append(None)
176 # --- Encoding ---
177 if isinstance(encoding, Encoding):
178 # user wants custom strategy? do it!
179 self._enc = encoding
180 else:
181 # use hammming encoding by default
182 self._enc = Encoding("hamming", encoding)
184 if self._enc.is_golomb:
185 self._enc._n_qubits = n_qubits
187 # Number of possible inputs
188 self.n_input_feat = len(self._enc)
189 log.debug(f"Number of input features: {self.n_input_feat}")
191 # Trainable frequencies, default initialization as in arXiv:2309.03279v2
192 self.enc_params = jnp.ones((self.n_layers, self.n_qubits, self.n_input_feat))
194 # Per-feature pulse-parameter sizes/offsets used to slice
195 # enc_pulse_params in _iec under "all_pulse" mode. Only encodings whose
196 # gates all have a pulse parametrization are supported (golomb and
197 # custom callables do not).
198 # TODO: golomb should be doable but needs a closer investigation
199 self._enc_pulse_sizes: List[int] = []
200 self._enc_pulse_capable = not self._enc.is_golomb
201 if self._enc_pulse_capable:
202 for g in self._enc._gates:
203 if pinfo.gate_by_name(g) is None:
204 self._enc_pulse_capable = False
205 self._enc_pulse_sizes = []
206 break
207 self._enc_pulse_sizes.append(pinfo.gate_by_name(g).size)
209 self._enc_pulse_offsets: List[int] = list(
210 np.cumsum([0, *self._enc_pulse_sizes[:-1]])
211 )
212 self._enc_pulse_shape: Tuple[int, int, int] = (
213 self.n_layers,
214 self.n_qubits,
215 sum(self._enc_pulse_sizes),
216 )
218 # --- Data-Reuploading ---
220 # Keep as NumPy array (not JAX) so that ``if data_reupload[q, idx]``
221 # in _iec remains a concrete Python bool even under jax.jit tracing.
222 # note that setting this will also update self.degree and self.frequencies
223 # and in consequence also self.has_dru
224 self.data_reupload = data_reupload
226 # check for the highest degree among all input dimensions
227 if self.has_dru:
228 impl_n_layers: int = n_layers + 1 # we need L+1 according to Schuld et al.
229 else:
230 impl_n_layers = n_layers
231 log.info(f"Number of implicit layers: {impl_n_layers}.")
233 # --- Ansatz ---
234 # only weak check for str. We trust the user to provide sth useful
235 if isinstance(circuit_type, str):
236 self.pqc: Callable[[Optional[jnp.ndarray], int], int] = getattr(
237 Ansaetze, circuit_type or "No_Ansatz"
238 )()
239 else:
240 self.pqc = circuit_type()
241 log.info(f"Using Ansatz {circuit_type}.")
243 # calculate the shape of the parameter vector here, we will re-use this in init.
244 params_per_layer = self.pqc.n_params_per_layer(self.n_qubits)
245 self._params_shape: Tuple[int, int] = (impl_n_layers, params_per_layer)
246 log.info(f"Parameters per layer: {params_per_layer}")
248 pulse_params_per_layer = self.pqc.n_pulse_params_per_layer(self.n_qubits)
249 self._pulse_params_shape: Tuple[int, int] = (
250 impl_n_layers,
251 pulse_params_per_layer,
252 )
254 # intialize to None as we can't know this yet
255 self._batch_shape = None
257 # this will also be re-used in the init method,
258 # however, only if nothing is provided
259 self._inialization_strategy = initialization
260 self._initialization_domain = initialization_domain
262 # ..here! where we only require a JAX random key
263 self.random_key = self.initialize_params(random.key(random_seed))
265 # Initializing pulse params
266 self.pulse_params: jnp.ndarray = jnp.ones((1, *self._pulse_params_shape))
268 log.info(f"Initialized pulse parameters with shape {self.pulse_params.shape}.")
270 # Initializing encoding pulse params (element-wise scalers, ones by
271 # default). Batch-first convention, mirroring pulse_params.
272 self.enc_pulse_params: jnp.ndarray = jnp.ones((1, *self._enc_pulse_shape))
274 log.info(
275 f"Initialized encoding pulse parameters with shape "
276 f"{self.enc_pulse_params.shape}."
277 )
279 # Initialise the jaqsi Script that wraps _variational.
280 # No device selection needed - jaqsi auto-routes between statevector
281 # and density-matrix simulation based on whether noise channels are
282 # present on the tape.
283 self.script = js.Script(f=self._variational, n_qubits=self.n_qubits)
285 @property
286 def noise_params(self) -> Optional[Dict[str, Union[float, Dict[str, float]]]]:
287 """
288 Gets the noise parameters of the model.
290 Returns:
291 Optional[Dict[str, float]]: A dictionary of
292 noise parameters or None if not set.
293 """
294 return self._noise_params
296 @noise_params.setter
297 def noise_params(
298 self, kvs: Optional[Dict[str, Union[float, Dict[str, float]]]]
299 ) -> None:
300 """
301 Sets the noise parameters of the model.
303 Typically a "noise parameter" refers to the error probability.
304 ThermalRelaxation is a special case, and supports a dict as value with
305 structure:
306 "ThermalRelaxation":
307 {
308 "t1": 2000, # relative t1 time.
309 "t2": 1000, # relative t2 time
310 "t_factor" 1: # relative gate time factor
311 },
313 Args:
314 kvs (Optional[Dict[str, Union[float, Dict[str, float]]]]): A
315 dictionary of noise parameters. If all values are 0.0, the noise
316 parameters are set to None.
318 Returns:
319 None
320 """
321 self._noise_params = self._normalize_noise_params(kvs)
323 @staticmethod
324 def _normalize_noise_params(
325 kvs: Optional[Dict[str, Union[float, Dict[str, float]]]],
326 ) -> Optional[Dict[str, Union[float, Dict[str, float]]]]:
327 """
328 Fill in defaults and validate a noise parameter dictionary.
330 Args:
331 kvs (Optional[Dict[str, Union[float, Dict[str, float]]]]): A
332 dictionary of noise parameters.
334 Returns:
335 Optional[Dict[str, Union[float, Dict[str, float]]]]: The normalized
336 dictionary, or None if all values are 0.0.
337 """
338 # set to None if only zero values provided
339 if kvs is not None and all(v == 0.0 for v in kvs.values()):
340 kvs = None
342 # set default values
343 if kvs is not None:
344 defaults = {
345 "BitFlip": 0.0,
346 "PhaseFlip": 0.0,
347 "Depolarizing": 0.0,
348 "MultiQubitDepolarizing": 0.0,
349 "AmplitudeDamping": 0.0,
350 "PhaseDamping": 0.0,
351 "GateError": 0.0,
352 "ThermalRelaxation": None,
353 "StatePreparation": 0.0,
354 "Measurement": 0.0,
355 }
356 for key, default_val in defaults.items():
357 kvs.setdefault(key, default_val)
359 # check if there are any keys not supported
360 for key in kvs.keys():
361 if key not in defaults:
362 warnings.warn(
363 f"Noise type {key} is not supported by this package",
364 UserWarning,
365 )
367 # check valid params for thermal relaxation noise channel
368 tr_params = kvs["ThermalRelaxation"]
369 if isinstance(tr_params, dict):
370 tr_params.setdefault("t1", 0.0)
371 tr_params.setdefault("t2", 0.0)
372 tr_params.setdefault("t_factor", 0.0)
373 valid_tr_keys = {"t1", "t2", "t_factor"}
374 for k in tr_params.keys():
375 if k not in valid_tr_keys:
376 warnings.warn(
377 f"Thermal Relaxation parameter {k} is not supported "
378 f"by this package",
379 UserWarning,
380 )
381 if not all(tr_params.values()) or tr_params["t2"] > 2 * tr_params["t1"]:
382 warnings.warn(
383 "Received invalid values for Thermal Relaxation noise "
384 "parameter. Thermal relaxation is not applied!",
385 UserWarning,
386 )
387 kvs["ThermalRelaxation"] = 0.0
389 return kvs
391 @property
392 def output_qubit(self) -> List[int]:
393 """Deprecated alias for :attr:`observables`; returns the measured wires."""
394 warnings.warn(
395 "output_qubit is deprecated, use observables instead.",
396 DeprecationWarning,
397 stacklevel=2,
398 )
399 return self._measured_wires
401 @output_qubit.setter
402 def output_qubit(self, value: Union[int, List[int]]) -> None:
403 warnings.warn(
404 "output_qubit is deprecated, use observables instead.",
405 DeprecationWarning,
406 stacklevel=2,
407 )
408 self.observables = value
410 @property
411 def observables(self) -> List:
412 """The custom :class:`~jaqsi.operations.Operation` observables,
413 or the list of measured wires when using the default PauliZ readout.
415 With a list of observables, ``__call__`` and ``execution_type="expval"``
416 returns one expectation value per observable instead of one ``PauliZ``
417 per measured qubit.
418 """
419 return (
420 self._observables if self._observables is not None else self._measured_wires
421 )
423 @observables.setter
424 def observables(self, value: Union[int, List, None]) -> None:
425 if value is None:
426 self._observables = None
427 self._measured_wires = list(range(self.n_qubits))
428 elif (
429 isinstance(value, list)
430 and value
431 and all(isinstance(o, op.Operation) for o in value)
432 ):
433 self._observables = list(value)
434 self._measured_wires = list(range(self.n_qubits))
435 elif isinstance(value, list) and any(
436 isinstance(o, op.Operation) for o in value
437 ):
438 raise ValueError(
439 "observables list must contain either qubit indices or "
440 "Operation objects, not a mix."
441 )
442 else:
443 # qubit specification: normalize into the measured wire list
444 self._observables = None
445 if isinstance(value, list):
446 assert len(value) <= self.n_qubits, (
447 f"Size of observables {len(value)} cannot be larger than "
448 f"number of qubits {self.n_qubits}."
449 )
450 self._measured_wires = value
451 elif isinstance(value, int):
452 if value == -1:
453 self._measured_wires = list(range(self.n_qubits))
454 else:
455 assert value < self.n_qubits, (
456 f"Output qubit {value} cannot be larger than {self.n_qubits}."
457 )
458 self._measured_wires = [value]
459 else:
460 self._measured_wires = value
462 # recompute the result shape for the (possibly new) observable count
463 if hasattr(self, "_execution_type"):
464 self.execution_type = self.execution_type
466 @property
467 def execution_type(self) -> str:
468 """
469 Gets the execution type of the model.
471 Returns:
472 str: The execution type, one of 'density', 'expval', or 'probs'.
473 """
474 return self._execution_type
476 @execution_type.setter
477 def execution_type(self, value: str) -> None:
478 self._result_shape = self._compute_result_shape(value)
480 if value == "state" and not self.all_qubit_measurement:
481 warnings.warn(
482 f"{value} measurement ignores the measured subsystem, which is "
483 f"{self._measured_wires}.",
484 UserWarning,
485 )
487 if value != "expval" and getattr(self, "_observables", None) is not None:
488 warnings.warn(
489 f"Custom observables are ignored for execution_type={value!r}.",
490 UserWarning,
491 )
493 if value == "probs" and self.shots is None:
494 warnings.warn(
495 "Setting execution_type to probs without specifying shots.",
496 UserWarning,
497 )
499 if value == "density" and self.shots is not None:
500 raise ValueError("Setting execution_type to density with shots not None.")
502 self._execution_type = value
504 def _compute_result_shape(self, execution_type: str) -> Tuple[int, ...]:
505 """
506 Derive the per-sample output shape for an execution type.
508 Args:
509 execution_type (str): One of "density", "expval", "probs", "state".
511 Returns:
512 Tuple[int, ...]: The output shape of a single sample.
514 Raises:
515 ValueError: If execution_type is not supported.
516 """
517 if execution_type == "density":
518 return (
519 2 ** len(self._measured_wires),
520 2 ** len(self._measured_wires),
521 )
522 elif execution_type == "expval":
523 # custom observables (if provided) fix the number of expectation
524 # values; otherwise one PauliZ (or Z-parity) per measured qubit.
525 if getattr(self, "_observables", None) is not None:
526 return (len(self._observables),)
527 else:
528 return (len(self._measured_wires),)
529 elif execution_type == "probs":
530 # in case this is a list of parities,
531 # each pair has 2^len(qubits) probabilities
532 return (
533 (2,) * len(self._measured_wires)
534 if isinstance(self._measured_wires, (Tuple, List))
535 else (2,)
536 )
537 elif execution_type == "state":
538 return (2 ** len(self._measured_wires),)
539 else:
540 raise ValueError(f"Invalid execution type: {execution_type}.")
542 @property
543 def shots(self) -> Optional[int]:
544 """
545 Gets the number of shots to use for the quantum device.
547 Returns:
548 Optional[int]: The number of shots.
549 """
550 return self._shots
552 @shots.setter
553 def shots(self, value: Optional[int]) -> None:
554 """
555 Sets the number of shots to use for the quantum device.
557 Args:
558 value (Optional[int]): The number of shots.
559 If an integer less than or equal to 0 is provided, it is set to None.
561 Returns:
562 None
563 """
564 if type(value) is int and value <= 0:
565 value = None
566 self._shots = value
568 @property
569 def params(self) -> jnp.ndarray:
570 """Get the variational parameters of the model."""
571 return self._params
573 @params.setter
574 def params(self, value: jnp.ndarray) -> None:
575 """Set the variational parameters, ensuring batch dimension exists."""
576 if len(value.shape) == 2:
577 value = value.reshape(1, *value.shape)
579 self._params = value
581 @property
582 def enc_params(self) -> jnp.ndarray:
583 """Get the encoding parameters used for input transformation."""
584 return self._enc_params
586 @enc_params.setter
587 def enc_params(self, value: jnp.ndarray) -> None:
588 """Set the encoding parameters."""
589 self._enc_params = value
591 @property
592 def pulse_params(self) -> jnp.ndarray:
593 """Get the pulse parameters for pulse-mode gate execution."""
594 return self._pulse_params
596 @pulse_params.setter
597 def pulse_params(self, value: jnp.ndarray) -> None:
598 """Set the pulse parameters."""
599 self._pulse_params = value
601 @property
602 def enc_pulse_params(self) -> jnp.ndarray:
603 """Get the encoding pulse parameters for all_pulse-mode execution."""
604 return self._enc_pulse_params
606 @enc_pulse_params.setter
607 def enc_pulse_params(self, value: jnp.ndarray) -> None:
608 """Set the encoding pulse parameters."""
609 self._enc_pulse_params = value
611 @property
612 def data_reupload(self) -> np.ndarray:
613 """Get the data reupload mask."""
614 return self._data_reupload
616 @data_reupload.setter
617 def data_reupload(
618 self,
619 value: Union[bool, List[List[bool]], List[List[List[bool]]], np.ndarray],
620 ) -> None:
621 """Set the data reupload mask.
623 Always converts to a concrete NumPy boolean array so that
624 ``if data_reupload[q, idx]`` in :meth:`_iec` remains a plain
625 Python ``bool`` even inside JAX-traced functions (jit / grad / vmap).
626 """
627 # Process data reuploading strategy and set degree
628 if not isinstance(value, bool):
629 if not isinstance(value, np.ndarray):
630 value = np.array(value)
632 if len(value.shape) == 2:
633 assert value.shape == (
634 self.n_layers,
635 self.n_qubits,
636 ), (
637 f"Data reuploading array has wrong shape. \
638 Expected {(self.n_layers, self.n_qubits)} or\
639 {(self.n_layers, self.n_qubits, self.n_input_feat)},\
640 got {value.shape}."
641 )
642 value = value.reshape(*value.shape, 1)
643 value = np.repeat(value, self.n_input_feat, axis=2)
645 assert value.shape == (
646 self.n_layers,
647 self.n_qubits,
648 self.n_input_feat,
649 ), (
650 f"Data reuploading array has wrong shape. \
651 Expected {(self.n_layers, self.n_qubits, self.n_input_feat)},\
652 got {value.shape}."
653 )
655 log.debug(f"Data reuploading array:\n{value}")
656 else:
657 if value:
658 value = np.ones((self.n_layers, self.n_qubits, self.n_input_feat))
659 log.debug("Full data reuploading.")
660 else:
661 value = np.zeros((self.n_layers, self.n_qubits, self.n_input_feat))
662 value[0][0] = 1
663 log.debug("No data reuploading.")
665 # convert to boolean values
666 self._data_reupload = np.asarray(value).astype(bool)
668 self.degree: Tuple = tuple(
669 self._enc.get_n_freqs(self.data_reupload[..., i])
670 for i in range(self.n_input_feat)
671 )
673 self.frequencies: Tuple = tuple(
674 self._enc.get_spectrum(self.data_reupload[..., i])
675 for i in range(self.n_input_feat)
676 )
678 # Cache has_dru as a plain Python bool so that it can be used in
679 # Python ``if`` statements even inside JAX-traced functions.
680 self._has_dru: bool = bool(max(int(np.max(f)) for f in self._frequencies) > 1)
682 @property
683 def degree(self) -> Tuple:
684 """Get the degree of the model."""
685 return self._degree
687 @degree.setter
688 def degree(self, value: Tuple):
689 self._degree = value
691 @property
692 def frequencies(self) -> Tuple:
693 """Get the frequencies of the model."""
694 return self._frequencies
696 @frequencies.setter
697 def frequencies(self, value: Tuple):
698 self._frequencies = value
700 def exact_spectrum(self, method: str = "tree") -> Tuple[np.ndarray, ...]:
701 """Compute the exact per-feature Fourier spectrum via the FourierTree.
703 Unlike :attr:`frequencies` -- a naive per-feature estimate derived purely
704 from the encoding, which can *overestimate* the spectrum (some
705 coefficients are constrained to zero for all parameters) -- this builds
706 the analytical Fourier tree (Nemkov et al.) and returns, for each input
707 feature, the integer frequencies whose Fourier coefficient is not
708 identically zero. The result is always a subset of :attr:`frequencies`.
710 The support is derived purely symbolically (no parameter sampling): see
711 :meth:`~qml_essentials.coefficients.FourierTree.get_exact_support`.
712 With ``method="tree"`` (default), frequencies whose contributions cancel
713 identically across tree paths (e.g. two consecutive encodings combining
714 into a single rotation) are excluded exactly; this enumerates the
715 explicit tree, which can be infeasible for deep entangling circuits.
716 With ``method="dp"``, a merged-state dynamic program derives the support
717 without enumerating paths, which scales to deep circuits at the cost of
718 not detecting identical cross-path cancellations.
720 Requires a Clifford + Pauli-rotation ansatz (see
721 :class:`~qml_essentials.pauli.PauliCircuit`); other gate sets raise
722 ``NotImplementedError`` during tree construction.
724 Args:
725 method (str): ``"tree"`` (fully exact) or ``"dp"`` (scalable).
727 Returns:
728 Tuple[np.ndarray, ...]: One sorted integer frequency array per input
729 feature (same layout as :attr:`frequencies`).
730 """
731 from qml_essentials.coefficients import FourierTree # avoid circular imp.
733 tree = FourierTree(self)
735 # Position of each model feature within the tree's frequency vectors.
736 feature_pos = {feat: i for i, feat in enumerate(tree.features)}
738 # Union of the symbolic supports over all observables (roots).
739 support = set()
740 for freqs in tree.get_exact_support(method=method):
741 farr = np.asarray(freqs)
742 for k in range(farr.shape[0]):
743 key = (
744 (int(farr[k]),)
745 if farr.ndim == 1
746 else tuple(int(v) for v in farr[k])
747 )
748 support.add(key)
750 spectrum = []
751 for feat in range(self.n_input_feat):
752 if support and feat in feature_pos:
753 pos = feature_pos[feat]
754 vals = sorted({k[pos] for k in support})
755 else:
756 vals = [0]
757 spectrum.append(np.array(vals, dtype=int))
758 return tuple(spectrum)
760 @property
761 def has_dru(self) -> bool:
762 """Check if the model has data reupload."""
763 return self._has_dru
765 @property
766 def all_qubit_measurement(self) -> bool:
767 """Check if measurement is performed on all qubits."""
768 return self._measured_wires == list(range(self.n_qubits))
770 @property
771 def batch_shape(self) -> Tuple[int, ...]:
772 """
773 Get the batch shape (B_I, B_P, B_R, B_E).
774 If the model was not called before,
775 it returns (1, 1, 1, 1).
777 Returns:
778 Tuple[int, ...]: Tuple of (input_batch, param_batch, pulse_batch,
779 enc_pulse_batch). Returns (1, 1, 1, 1) if model has not been
780 called yet.
781 """
782 if self._batch_shape is None:
783 log.debug("Model was not called yet. Returning (1,1,1,1) as batch shape.")
784 return (1, 1, 1, 1)
785 return self._batch_shape
787 @property
788 def eff_batch_shape(self) -> Tuple[int, ...]:
789 """
790 Get the effective batch shape after applying repeat_batch_axis mask.
792 Returns:
793 Tuple[int, ...]: Effective batch dimensions, excluding zeros.
794 """
795 return self._eff_batch_shape_of(self.batch_shape)
797 def _eff_batch_shape_of(self, batch_shape: Tuple[int, ...]) -> Tuple[int, ...]:
798 """
799 Apply the repeat_batch_axis mask to a given batch shape.
801 Args:
802 batch_shape (Tuple[int, ...]): Batch shape (B_I, B_P, B_R, B_E).
804 Returns:
805 Tuple[int, ...]: Effective batch dimensions, excluding zeros.
806 """
807 batch_shape = np.array(batch_shape) * self.repeat_batch_axis
808 return batch_shape[batch_shape != 0]
810 def initialize_params(
811 self,
812 random_key: Optional[random.PRNGKey] = None,
813 repeat: int = 1,
814 initialization: Optional[str] = None,
815 initialization_domain: Optional[List[float]] = None,
816 ) -> random.PRNGKey:
817 """
818 Initialize the variational parameters of the model.
820 Args:
821 random_key (Optional[random.PRNGKey]): JAX random key for initialization.
822 If None, uses the model's internal random key.
823 repeat (int): Number of parameter sets to create (batch dimension).
824 Defaults to 1.
825 initialization (Optional[str]): Strategy for parameter initialization.
826 Options: "random", "zeros", "pi", "zero-controlled", "pi-controlled".
827 If None, uses the strategy specified in the constructor.
828 initialization_domain (Optional[List[float]]): Domain [min, max] for
829 random initialization. If None, uses the domain from constructor.
831 Returns:
832 random.PRNGKey: Updated random key after initialization.
834 Raises:
835 Exception: If an invalid initialization method is specified.
836 """
837 # Initializing params
838 params_shape = (repeat, *self._params_shape)
840 # use existing strategy if not specified
841 initialization = initialization or self._inialization_strategy
842 initialization_domain = initialization_domain or self._initialization_domain
844 random_key, sub_key = safe_random_split(
845 random_key if random_key is not None else self.random_key
846 )
848 def set_control_params(params: jnp.ndarray, value: float) -> jnp.ndarray:
849 indices = self.pqc.get_control_indices(self.n_qubits)
850 if indices is None:
851 warnings.warn(
852 f"Specified {initialization} but circuit\
853 does not contain controlled rotation gates.\
854 Parameters are intialized randomly.",
855 UserWarning,
856 )
857 else:
858 np_params = np.array(params)
859 np_params[:, :, indices[0] : indices[1] : indices[2]] = (
860 np.ones_like(params[:, :, indices[0] : indices[1] : indices[2]])
861 * value
862 )
863 params = jnp.array(np_params)
864 return params
866 if initialization == "random":
867 self.params: jnp.ndarray = random.uniform(
868 sub_key,
869 params_shape,
870 minval=initialization_domain[0],
871 maxval=initialization_domain[1],
872 )
873 elif initialization == "zeros":
874 self.params: jnp.ndarray = jnp.zeros(params_shape)
875 elif initialization == "pi":
876 self.params: jnp.ndarray = jnp.ones(params_shape) * jnp.pi
877 elif initialization == "zero-controlled":
878 self.params: jnp.ndarray = random.uniform(
879 sub_key,
880 params_shape,
881 minval=initialization_domain[0],
882 maxval=initialization_domain[1],
883 )
884 self.params = set_control_params(self.params, 0)
885 elif initialization == "pi-controlled":
886 self.params: jnp.ndarray = random.uniform(
887 sub_key,
888 params_shape,
889 minval=initialization_domain[0],
890 maxval=initialization_domain[1],
891 )
892 self.params = set_control_params(self.params, jnp.pi)
893 else:
894 raise Exception("Invalid initialization method")
896 log.info(
897 f"Initialized parameters with shape {self.params.shape}\
898 using strategy {initialization}."
899 )
901 return random_key
903 def next_key(self) -> random.PRNGKey:
904 """
905 Advance the internal random key and return a fresh sub key.
907 Intended for stochastic execution inside a JAX transform: a jitted
908 call is traced once and replays the key that was current at trace
909 time, so fresh randomness has to enter as an argument. Call this
910 outside the transform and pass the result as ``random_key``. Since the
911 key is an argument rather than a constant, this does not trigger
912 recompilation.
914 Returns:
915 random.PRNGKey: Fresh sub key, split off the internal key.
916 """
917 self.random_key, sub_key = safe_random_split(self.random_key)
918 return sub_key
920 def transform_input(
921 self, inputs: jnp.ndarray, enc_params: jnp.ndarray
922 ) -> jnp.ndarray:
923 """
924 Transform input data by scaling with encoding parameters.
926 Implements the input transformation as described in arXiv:2309.03279v2,
927 where inputs are linearly scaled by encoding parameters before being
928 used in the quantum circuit.
930 Args:
931 inputs (jnp.ndarray): Input data point of shape (n_input_feat,) or
932 (batch_size, n_input_feat).
933 enc_params (jnp.ndarray): Encoding weight scalar or vector used to
934 scale the input.
936 Returns:
937 jnp.ndarray: Transformed input, element-wise product of inputs
938 and enc_params.
939 """
940 return inputs * enc_params
942 def _iec(
943 self,
944 inputs: jnp.ndarray,
945 data_reupload: np.ndarray,
946 enc: Encoding,
947 enc_params: jnp.ndarray,
948 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
949 random_key: Optional[random.PRNGKey] = None,
950 enc_pulse_params: Optional[jnp.ndarray] = None,
951 pulse: bool = False,
952 ) -> None:
953 """
954 Apply Input Encoding Circuit (IEC) with angle encoding.
956 Encodes classical input data into the quantum circuit using rotation
957 gates (e.g., RX, RY, RZ). Supports data re-uploading at specified
958 positions in the circuit.
960 For Golomb encoding, a single multi-qubit diagonal unitary is applied
961 to all qubits simultaneously instead of per-qubit rotation gates.
963 Args:
964 inputs (jnp.ndarray): Input data of shape (n_input_feat,) or
965 (batch_size, n_input_feat).
966 data_reupload (np.ndarray): Boolean array of shape (n_qubits, n_input_feat)
967 indicating where to apply encoding gates.
968 enc (Encoding): Encoding strategy containing the encoding gate functions.
969 enc_params (jnp.ndarray): Encoding parameters of shape
970 (n_qubits, n_input_feat) used to scale inputs.
971 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
972 Noise parameters for gate-level noise simulation. Defaults to None.
973 random_key (Optional[random.PRNGKey]): JAX random key for stochastic
974 noise. Defaults to None.
975 enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
976 scalers of shape (n_qubits, n_enc_pulse_per_qubit) for the
977 current layer. Used when the encoding gates run at pulse level,
978 i.e. the model-level mode is "enc_pulse" or "all_pulse".
979 Defaults to None.
980 pulse (bool): Whether the encoding gates run at pulse level.
981 This is the backend selected for the encoding group, distinct
982 from the model-level modes (unitary, ansatz_pulse, enc_pulse,
983 all_pulse). Defaults to False.
985 Returns:
986 None: Gates are applied in-place to the quantum circuit.
987 """
988 # --- Golomb encoding: single multi-qubit gate on all qubits --------
989 if enc.is_golomb:
990 idx = 0 # Golomb encoding supports a single input feature
991 # Check if any qubit has re-uploading enabled for this layer
992 if data_reupload[:, idx].any():
993 random_key, sub_key = safe_random_split(random_key)
994 # Use the mean of enc_params across qubits as scalar scaling
995 # (Golomb acts on all qubits jointly)
996 mean_enc_param = jnp.mean(enc_params[:, idx])
997 all_wires = list(range(self.n_qubits))
998 enc[idx](
999 self.transform_input(inputs[..., idx], mean_enc_param),
1000 wires=all_wires,
1001 noise_params=noise_params,
1002 random_key=sub_key,
1003 )
1004 return
1006 # --- Standard per-qubit encoding -----------------------------------
1007 for q in range(self.n_qubits):
1008 # use the last dimension of the inputs (feature dimension)
1009 for idx in range(inputs.shape[-1]):
1010 if data_reupload[q, idx]:
1011 random_key, sub_key = safe_random_split(random_key)
1012 # TODO: consider merging this with the pulses.py manager
1013 pulse_kwargs = {}
1014 if pulse:
1015 # scale the calibrated pulse params by this gate's
1016 # scalers, as the pulse manager does for the ansatz
1017 off = self._enc_pulse_offsets[idx]
1018 size = self._enc_pulse_sizes[idx]
1019 base = pinfo.gate_by_name(enc._gates[idx]).params
1020 pulse_kwargs = dict(
1021 pulse_params=base * enc_pulse_params[q, off : off + size],
1022 pulse=True,
1023 )
1025 # use elipsis to index only the last dimension
1026 # as inputs are generally *not* qubit dependent
1027 enc[idx](
1028 self.transform_input(inputs[..., idx], enc_params[q, idx]),
1029 wires=q,
1030 noise_params=noise_params,
1031 random_key=sub_key,
1032 **pulse_kwargs,
1033 )
1035 @staticmethod
1036 def _debatch(value: jnp.ndarray, ndim: int) -> jnp.ndarray:
1037 """
1038 Drop a leading singleton batch axis (batch-first convention).
1040 Args:
1041 value (jnp.ndarray): Array to de-batch.
1042 ndim (int): Rank of a single (un-batched) element.
1044 Returns:
1045 jnp.ndarray: The array without its leading axis if that axis is a
1046 singleton batch dimension, otherwise the array unchanged.
1047 """
1048 if len(value.shape) > ndim and value.shape[0] == 1:
1049 return value[0]
1050 return value
1052 def _self_fallback(self, value: Any, name: str, warn: bool) -> Any:
1053 """
1054 Fall back to the model's own attribute when a parameter is not given.
1056 Args:
1057 value (Any): The provided value, or None.
1058 name (str): Name of the attribute to fall back to.
1059 warn (bool): Whether to warn when the fallback is used.
1061 Returns:
1062 Any: The provided value, or ``self.<name>`` if value is None.
1063 """
1064 if value is not None:
1065 return value
1066 if warn:
1067 warnings.warn(
1068 "Explicit call to `_circuit` or `_variational` detected: "
1069 f"`{name}` is None, using `self.{name}` instead.",
1070 RuntimeWarning,
1071 )
1072 return getattr(self, name)
1074 def _variational(
1075 self,
1076 params: jnp.ndarray,
1077 inputs: jnp.ndarray,
1078 pulse_params: Optional[jnp.ndarray] = None,
1079 random_key: Optional[random.PRNGKey] = None,
1080 enc_params: Optional[jnp.ndarray] = None,
1081 enc_pulse_params: Optional[jnp.ndarray] = None,
1082 gate_mode: str = "unitary",
1083 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
1084 ) -> None:
1085 """
1086 Build the variational quantum circuit structure.
1088 Constructs the circuit by applying state preparation, alternating
1089 variational ansatz layers with input encoding layers, and optional
1090 noise channels.
1092 The first six parameters (after ``self``) - ``params``, ``inputs``,
1093 ``pulse_params``, ``random_key``, ``enc_params``, ``enc_pulse_params`` -
1094 are the batchable positional arguments.
1095 The remaining keyword arguments are broadcast across the batch.
1097 Args:
1098 params (jnp.ndarray): Variational parameters of shape
1099 (n_layers, n_params_per_layer).
1100 inputs (jnp.ndarray): Input data of shape (n_input_feat,).
1101 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers of shape
1102 (n_layers, n_pulse_params_per_layer) for pulse-mode execution.
1103 Defaults to None (uses model's pulse_params).
1104 random_key (Optional[random.PRNGKey]): JAX random key for stochastic
1105 operations. Defaults to None.
1106 enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
1107 (n_qubits, n_input_feat). Defaults to None (uses model's enc_params).
1108 enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
1109 scalers of shape (n_layers, n_qubits, n_enc_pulse_per_qubit) for
1110 "all_pulse" execution. Defaults to None (uses model's
1111 enc_pulse_params).
1112 gate_mode (str): Gate execution mode, one of "unitary",
1113 "ansatz_pulse", "enc_pulse" or "all_pulse". "ansatz_pulse" runs
1114 the ansatz and state preparation as pulses (encoding stays
1115 unitary); "enc_pulse" runs only the encoding gates as pulses;
1116 "all_pulse" runs both as pulses. Defaults to "unitary".
1117 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
1118 Noise parameters for simulation. Defaults to None.
1120 Returns:
1121 None: Gates are applied in-place to the quantum circuit.
1123 Note:
1124 Issues RuntimeWarning if called directly without providing parameters
1125 that would normally be passed through the forward method.
1126 """
1127 # which backend the ansatz / state-prep gates and the encoding gates use
1128 use_pulse, enc_use_pulse = GATE_MODES[gate_mode]
1130 # TODO: rework and double check params shape
1131 params = self._debatch(params, 2)
1132 inputs = self._debatch(inputs, 1)
1134 # TODO: Raise warning if trainable frequencies is True, or similar. I.e., no
1135 # warning if user does not care for frequencies or enc_params
1136 enc_params = self._self_fallback(
1137 enc_params, "enc_params", self.trainable_frequencies
1138 )
1140 pulse_params = self._self_fallback(pulse_params, "pulse_params", use_pulse)
1141 pulse_params = self._debatch(pulse_params, 2)
1143 enc_pulse_params = self._self_fallback(
1144 enc_pulse_params, "enc_pulse_params", enc_use_pulse
1145 )
1146 enc_pulse_params = self._debatch(enc_pulse_params, 3)
1148 noise_params = self._self_fallback(
1149 noise_params, "noise_params", self.noise_params is not None
1150 )
1152 if noise_params is not None:
1153 random_key = self._self_fallback(random_key, "random_key", True)
1154 self._apply_state_prep_noise(noise_params=noise_params)
1156 # state preparation
1157 for q in range(self.n_qubits):
1158 for _sp, sp_pulse_params in zip(self._sp, self.sp_pulse_params):
1159 random_key, sub_key = safe_random_split(random_key)
1160 _sp(
1161 wires=q,
1162 pulse_params=sp_pulse_params,
1163 noise_params=noise_params,
1164 random_key=sub_key,
1165 pulse=use_pulse,
1166 )
1168 # circuit building
1169 for layer in range(0, self.n_layers):
1170 random_key, sub_key = safe_random_split(random_key)
1171 # ansatz layers
1172 self.pqc(
1173 params[layer],
1174 self.n_qubits,
1175 pulse_params=pulse_params[layer],
1176 noise_params=noise_params,
1177 random_key=sub_key,
1178 pulse=use_pulse,
1179 )
1181 random_key, sub_key = safe_random_split(random_key)
1182 # encoding layers
1183 self._iec(
1184 inputs,
1185 data_reupload=self.data_reupload[layer],
1186 enc=self._enc,
1187 enc_params=enc_params[layer],
1188 noise_params=noise_params,
1189 random_key=sub_key,
1190 enc_pulse_params=enc_pulse_params[layer],
1191 pulse=enc_use_pulse,
1192 )
1194 # final ansatz layer
1195 if self.has_dru: # same check as in init
1196 random_key, sub_key = safe_random_split(random_key)
1197 self.pqc(
1198 params[self.n_layers],
1199 self.n_qubits,
1200 pulse_params=pulse_params[-1],
1201 noise_params=noise_params,
1202 random_key=sub_key,
1203 pulse=use_pulse,
1204 )
1206 # channel noise
1207 if noise_params is not None:
1208 self._apply_general_noise(noise_params=noise_params)
1210 def _build_obs(
1211 self, execution_type: Optional[str] = None
1212 ) -> Tuple[str, List[op.Operation]]:
1213 """Build the jaqsi measurement type and observable list.
1215 Translates the model's ``execution_type`` and ``observables``
1216 settings into parameters suitable for
1217 :meth:`~jaqsi.Script.execute`.
1219 Args:
1220 execution_type: Measurement type to build for. If ``None``, the
1221 model's current ``execution_type`` is used.
1223 Returns:
1224 Tuple ``(meas_type, obs)`` where *meas_type* is one of
1225 ``"expval"``, ``"probs"``, ``"density"``, ``"state"`` and *obs*
1226 is a (possibly empty) list of :class:`Operation` observables.
1227 """
1228 if execution_type is None:
1229 execution_type = self.execution_type
1231 if execution_type == "density":
1232 return "density", []
1234 if execution_type == "state":
1235 return "state", []
1237 if execution_type == "expval":
1238 if self._observables is not None:
1239 return "expval", list(self._observables)
1240 obs: List[op.Operation] = []
1241 for qubit_spec in self._measured_wires:
1242 if isinstance(qubit_spec, int):
1243 obs.append(gateset.PauliZ(wires=qubit_spec))
1244 else:
1245 # parity: Z \\otimes Z \\otimes …
1246 obs.append(js.build_parity_observable(list(qubit_spec)))
1247 return "expval", obs
1249 if execution_type == "probs":
1250 # probs are computed on the full system; subsystem
1251 # marginalisation is handled in _postprocess_res
1252 return "probs", []
1254 raise ValueError(f"Invalid execution_type: {execution_type}.")
1256 def _apply_state_prep_noise(
1257 self, noise_params: Dict[str, Union[float, Dict[str, float]]]
1258 ) -> None:
1259 """
1260 Apply state preparation noise to all qubits.
1262 Simulates imperfect state preparation by applying BitFlip errors
1263 to each qubit with the specified probability.
1265 Args:
1266 noise_params (Dict[str, Union[float, Dict[str, float]]]): Dictionary
1267 containing noise parameters. Uses the "StatePreparation" key
1268 for the BitFlip probability.
1270 Returns:
1271 None: Noise channels are applied in-place to the circuit.
1272 """
1273 p = noise_params.get("StatePreparation", 0.0)
1274 if p > 0:
1275 for q in range(self.n_qubits):
1276 noise.BitFlip(p, wires=q)
1278 def _apply_general_noise(
1279 self, noise_params: Dict[str, Union[float, Dict[str, float]]]
1280 ) -> None:
1281 """
1282 Apply general noise channels to all qubits.
1284 Applies various decoherence and error channels after the circuit
1285 execution, simulating environmental noise effects.
1287 Args:
1288 noise_params (Dict[str, Union[float, Dict[str, float]]]): Dictionary
1289 containing noise parameters with the following supported keys:
1290 - "AmplitudeDamping" (float): Probability for amplitude damping.
1291 - "PhaseDamping" (float): Probability for phase damping.
1292 - "Measurement" (float): Probability for measurement error (BitFlip).
1293 - "ThermalRelaxation" (Dict): Dictionary with keys "t1", "t2",
1294 "t_factor" for thermal relaxation simulation.
1296 Returns:
1297 None: Noise channels are applied in-place to the circuit.
1299 Note:
1300 Gate-level noise (e.g., GateError) is handled separately in the
1301 Gates.Noise module and applied at the individual gate level.
1302 """
1303 amp_damp = noise_params.get("AmplitudeDamping", 0.0)
1304 phase_damp = noise_params.get("PhaseDamping", 0.0)
1305 thermal_relax = noise_params.get("ThermalRelaxation", 0.0)
1306 meas = noise_params.get("Measurement", 0.0)
1307 for q in range(self.n_qubits):
1308 if amp_damp > 0:
1309 noise.AmplitudeDamping(amp_damp, wires=q)
1310 if phase_damp > 0:
1311 noise.PhaseDamping(phase_damp, wires=q)
1312 if meas > 0:
1313 noise.BitFlip(meas, wires=q)
1314 if isinstance(thermal_relax, dict):
1315 t1 = thermal_relax["t1"]
1316 t2 = thermal_relax["t2"]
1317 t_factor = thermal_relax["t_factor"]
1318 circuit_depth = self._get_circuit_depth()
1319 tg = circuit_depth * t_factor
1320 noise.ThermalRelaxationError(1.0, t1, t2, tg, q)
1322 def _get_circuit_depth(self, inputs: Optional[jnp.ndarray] = None) -> int:
1323 """
1324 Calculate the depth of the quantum circuit.
1326 Records the circuit onto a tape (without noise) and computes the
1327 depth as the length of the critical path: each gate is scheduled
1328 at the earliest time step after all of its qubits are free.
1330 Args:
1331 inputs (Optional[jnp.ndarray]): Input data for circuit evaluation.
1332 If None, default zero inputs are used.
1334 Returns:
1335 int: The circuit depth (longest path of gates in the circuit).
1336 """
1337 # Return cached value if available
1338 if hasattr(self, "_cached_circuit_depth"):
1339 return self._cached_circuit_depth
1341 inputs = self._inputs_validation(inputs)
1343 # Temporarily clear noise_params to prevent _variational from
1344 # picking them up (which would call _apply_general_noise ->
1345 # _get_circuit_depth again, causing infinite recursion).
1346 saved_noise = self._noise_params
1347 self._noise_params = None
1349 with recording() as tape:
1350 self._variational(
1351 self.params[0] if self.params.ndim == 3 else self.params,
1352 inputs[0] if inputs.ndim == 2 else inputs,
1353 noise_params=None,
1354 )
1356 self._noise_params = saved_noise
1358 # Filter out noise channels - only count unitary gates
1359 ops = [o for o in tape if not isinstance(o, KrausChannel)]
1361 if not ops:
1362 self._cached_circuit_depth = 0
1363 return 0
1365 # Schedule each gate at the earliest time step where all its wires
1366 # are free. ``wire_busy[q]`` tracks the next free time step for
1367 # qubit ``q``.
1368 wire_busy: Dict[int, int] = {}
1369 depth = 0
1370 for gate in ops:
1371 start = max((wire_busy.get(w, 0) for w in gate.wires), default=0)
1372 end = start + 1
1373 for w in gate.wires:
1374 wire_busy[w] = end
1375 depth = max(depth, end)
1377 self._cached_circuit_depth = depth
1378 return depth
1380 def draw(
1381 self,
1382 inputs: Optional[jnp.ndarray] = None,
1383 figure: str = "text",
1384 **kwargs: Any,
1385 ) -> Union[str, Any]:
1386 """Visualize the quantum circuit.
1388 Records the circuit tape (without noise) and renders the gate
1389 sequence using the requested backend.
1391 Args:
1392 inputs (Optional[jnp.ndarray]): Input data for the circuit.
1393 If ``None``, default zero inputs are used.
1394 figure (str): Rendering backend. One of:
1396 * ``"text"`` - ASCII art (returned as a ``str``).
1397 * ``"mpl"`` - Matplotlib figure (returns ``(fig, ax)``).
1398 * ``"tikz"`` - LaTeX/TikZ ``quantikz`` code (returns a
1399 :class:`TikzFigure`).
1400 * ``"pulse"`` - Pulse schedule (returns ``(fig, axes)``).
1401 Only meaningful for pulse-mode models.
1403 **kwargs: Extra options forwarded to the drawing backend
1404 (e.g. ``gate_values=True``).
1406 Returns:
1407 Depends on figure:
1409 * ``"text"`` -> ``str``
1410 * ``"mpl"`` -> ``(matplotlib.figure.Figure, matplotlib.axes.Axes)``
1411 * ``"tikz"`` -> :class:`TikzFigure`
1413 Raises:
1414 ValueError: If figure is not one of the supported modes.
1415 """
1416 inputs = self._inputs_validation(inputs)
1417 params = self.params[0] if self.params.ndim == 3 else self.params
1418 inp = inputs[0] if inputs.ndim == 2 else inputs
1420 if figure == "pulse":
1421 return self.draw_pulse(inputs=inputs, **kwargs)
1423 # Record without noise to get a clean circuit
1424 saved_noise = self._noise_params
1425 self._noise_params = None
1427 draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
1428 result = draw_script.draw(
1429 figure=figure,
1430 args=(params, inp),
1431 kwargs={"noise_params": None},
1432 **kwargs,
1433 )
1435 self._noise_params = saved_noise
1436 return result
1438 def draw_pulse(
1439 self,
1440 inputs: Optional[jnp.ndarray] = None,
1441 **kwargs: Any,
1442 ) -> Any:
1443 """Visualize the pulse schedule for the circuit.
1445 Records the circuit in pulse mode and collects PulseEvents
1446 automatically via the pulse-event tape, then renders them.
1448 State preparation, ansatz and encoding gates are all rendered as
1449 pulses. Encodings without a pulse parametrization (golomb and custom
1450 callables) are omitted from the schedule.
1452 Args:
1453 inputs: Input data. If ``None``, default zero inputs are used.
1454 **kwargs: Forwarded to
1455 :func:`~jaqsi.drawing.draw_pulse_schedule`
1456 (e.g. ``show_carrier=True``, ``n_samples=300``).
1458 Returns:
1459 ``(fig, axes)`` — Matplotlib Figure and array of Axes.
1460 """
1461 if "gate_mode" in kwargs:
1462 warnings.warn(
1463 "draw_pulse no longer takes gate_mode, every gate group with a "
1464 "pulse representation is drawn.",
1465 DeprecationWarning,
1466 stacklevel=2,
1467 )
1468 kwargs.pop("gate_mode")
1470 inputs = self._inputs_validation(inputs)
1471 params = self.params[0] if self.params.ndim == 3 else self.params
1472 inp = inputs[0] if inputs.ndim == 2 else inputs
1474 # pass the model's own pulse parameters, so that _variational does not
1475 # fall back to them with a warning. Both are batch-first, so drawing
1476 # picks the first set, same as params above
1477 record_kwargs: Dict[str, Any] = {
1478 "gate_mode": "all_pulse" if self._enc_pulse_capable else "ansatz_pulse",
1479 "noise_params": None,
1480 "pulse_params": self.pulse_params[0],
1481 }
1482 if self._enc_pulse_capable:
1483 record_kwargs["enc_pulse_params"] = self.enc_pulse_params[0]
1485 draw_script = js.Script(f=self._variational, n_qubits=self.n_qubits)
1486 return draw_script.draw(
1487 figure="pulse",
1488 args=(params, inp),
1489 kwargs=record_kwargs,
1490 **kwargs,
1491 )
1493 def __repr__(self) -> str:
1494 """Return text representation of the quantum circuit model."""
1495 return self.draw(figure="text")
1497 def __str__(self) -> str:
1498 """Return string representation of the quantum circuit model."""
1499 return self.draw(figure="text")
1501 def _params_validation(
1502 self, params: Optional[jnp.ndarray], stash: bool = True
1503 ) -> jnp.ndarray:
1504 """
1505 Validate and normalize variational parameters.
1507 Ensures parameters have the correct shape with a batch dimension,
1508 and updates the model's internal parameters if new ones are provided.
1510 Args:
1511 params (Optional[jnp.ndarray]): Variational parameters to validate.
1512 If None, returns the model's current parameters.
1513 stash (bool): Whether to store the validated parameters on the
1514 model. Set to False on the pure :meth:`apply` path, where
1515 stashing a JAX tracer would leak it across calls. Defaults
1516 to True.
1518 Returns:
1519 jnp.ndarray: Validated parameters with shape
1520 (batch_size, n_layers, n_params_per_layer).
1521 """
1522 # append batch axis if not provided
1523 if params is not None:
1524 if len(params.shape) == 2:
1525 # jnp (not np) so params stays a JAX array under autodiff /
1526 # jit; mirrors the pulse_params handling below.
1527 params = jnp.expand_dims(params, axis=0)
1529 # Avoid stashing JAX tracers on ``self``: under an outer
1530 # transform (e.g. ``jit``/``jacrev``) the tracer becomes invalid
1531 # once the transform returns, and a subsequent read of
1532 # ``self.params`` would feed a leaked tracer into the next
1533 # call (raising ``UnexpectedTracerError``).
1534 if stash and not isinstance(params, jax.core.Tracer):
1535 self.params = params
1536 elif stash:
1537 log.debug(
1538 "`params` is a JAX tracer; `self.params` is left at its "
1539 "previous value. Anything reading model state afterwards "
1540 "(draw, Entanglement, Expressibility, or a call that omits "
1541 "`params`) will see the stale parameters - assign "
1542 "`model.params` explicitly if you need the state to follow."
1543 )
1544 else:
1545 params = self.params
1547 return params
1549 def _pulse_params_validation(
1550 self, pulse_params: Optional[jnp.ndarray], stash: bool = True
1551 ) -> jnp.ndarray:
1552 """
1553 Validate and normalize pulse parameters.
1555 Ensures pulse parameters are set, using model defaults if not provided.
1557 Args:
1558 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers.
1559 If None, returns the model's current pulse parameters.
1560 stash (bool): Whether to store the validated parameters on the
1561 model. See :meth:`_params_validation`. Defaults to True.
1563 Returns:
1564 jnp.ndarray: Validated pulse parameters with shape
1565 (batch_size, n_layers, n_pulse_params_per_layer).
1566 """
1567 if pulse_params is None:
1568 pulse_params = self.pulse_params
1569 else:
1570 # ensure batch dimension exists (batch-first convention)
1571 if len(pulse_params.shape) == 2:
1572 pulse_params = jnp.expand_dims(pulse_params, axis=0)
1573 # See note in _params_validation: never stash JAX tracers on
1574 # ``self``.
1575 if stash and not isinstance(pulse_params, jax.core.Tracer):
1576 self.pulse_params = pulse_params
1577 elif stash:
1578 log.debug(
1579 "`pulse_params` is a JAX tracer; `self.pulse_params` is "
1580 "left at its previous value."
1581 )
1583 return pulse_params
1585 def _enc_pulse_params_validation(
1586 self, enc_pulse_params: Optional[jnp.ndarray], stash: bool = True
1587 ) -> jnp.ndarray:
1588 """
1589 Validate and normalize encoding pulse parameters.
1591 Ensures encoding pulse parameters are set (using model defaults if not
1592 provided) and carry a leading batch dimension.
1594 Args:
1595 enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
1596 scalers. If None, returns the model's current encoding pulse
1597 parameters.
1599 Returns:
1600 jnp.ndarray: Validated encoding pulse parameters with shape
1601 (batch_size, n_layers, n_qubits, n_enc_pulse_per_qubit).
1603 Raises:
1604 ValueError: If the trailing dimensions do not match the model's
1605 encoding pulse parameter shape.
1606 """
1607 if enc_pulse_params is None:
1608 enc_pulse_params = self.enc_pulse_params
1609 else:
1610 # ensure batch dimension exists (batch-first convention)
1611 if len(enc_pulse_params.shape) == 3:
1612 enc_pulse_params = jnp.expand_dims(enc_pulse_params, axis=0)
1613 if enc_pulse_params.shape[1:] != self._enc_pulse_shape:
1614 raise ValueError(
1615 f"enc_pulse_params trailing shape {enc_pulse_params.shape[1:]} "
1616 f"does not match expected {self._enc_pulse_shape}."
1617 )
1618 # See note in _params_validation: never stash JAX tracers on
1619 # ``self``.
1620 if stash and not isinstance(enc_pulse_params, jax.core.Tracer):
1621 self.enc_pulse_params = enc_pulse_params
1623 return enc_pulse_params
1625 def _resolve_gate_mode(
1626 self,
1627 gate_mode: Optional[str],
1628 pulse_params: Optional[jnp.ndarray],
1629 enc_pulse_params: Optional[jnp.ndarray],
1630 stacklevel: int = 3,
1631 ) -> str:
1632 """
1633 Determine which gate groups run at pulse level.
1635 The mode follows from the pulse parameters that were provided:
1636 ``pulse_params`` lowers the ansatz and state-preparation gates,
1637 ``enc_pulse_params`` lowers the encoding gates, both together lower
1638 everything.
1640 Args:
1641 gate_mode (Optional[str]): Deprecated explicit mode. If None, the
1642 mode is inferred from the pulse parameters.
1643 pulse_params (Optional[jnp.ndarray]): Ansatz pulse-parameter scalers.
1644 enc_pulse_params (Optional[jnp.ndarray]): Encoding pulse-parameter
1645 scalers.
1646 stacklevel (int): Frames between this method and the user call, so
1647 that the deprecation warning points at the call site. Defaults
1648 to 3, which is correct for a direct caller.
1650 Returns:
1651 str: One of the keys of ``GATE_MODES``.
1653 Raises:
1654 ValueError: If the encoding gates would run at pulse level but the
1655 encoding has no pulse parametrization, or if an explicitly
1656 passed (deprecated) gate_mode is unknown or inconsistent with
1657 the provided pulse parameters.
1658 """
1659 if gate_mode is None:
1660 if pulse_params is not None and enc_pulse_params is not None:
1661 gate_mode = "all_pulse"
1662 elif pulse_params is not None:
1663 gate_mode = "ansatz_pulse"
1664 elif enc_pulse_params is not None:
1665 gate_mode = "enc_pulse"
1666 else:
1667 gate_mode = "unitary"
1668 else:
1669 warnings.warn(
1670 "gate_mode is deprecated, the mode is inferred from the "
1671 "provided pulse parameters instead. Pass pulse_params to run "
1672 "the ansatz at pulse level and enc_pulse_params to run the "
1673 "encoding at pulse level, e.g. "
1674 "model(pulse_params=model.pulse_params).",
1675 DeprecationWarning,
1676 stacklevel=stacklevel,
1677 )
1678 # consistency checks, only reachable via the deprecated argument
1679 if gate_mode not in GATE_MODES:
1680 raise ValueError(
1681 f"Unknown gate_mode: {gate_mode}. Use one of {list(GATE_MODES)}."
1682 )
1683 if pulse_params is not None and gate_mode not in _ANSATZ_PULSE_MODES:
1684 raise ValueError(
1685 f"pulse_params were provided but gate_mode is not one of "
1686 f"{list(_ANSATZ_PULSE_MODES)}. Either switch gate_mode or do "
1687 "not pass pulse_params."
1688 )
1689 if enc_pulse_params is not None and gate_mode not in _ENC_PULSE_MODES:
1690 raise ValueError(
1691 f"enc_pulse_params were provided but gate_mode is not one of "
1692 f"{list(_ENC_PULSE_MODES)}. Either switch gate_mode or do not "
1693 "pass enc_pulse_params."
1694 )
1696 if gate_mode in _ENC_PULSE_MODES and not self._enc_pulse_capable:
1697 raise ValueError(
1698 "Pulse-level encoding requires an encoding whose gates have a "
1699 "pulse parametrization (golomb and custom callables do not). "
1700 "Do not pass enc_pulse_params for this model."
1701 )
1703 return gate_mode
1705 def _enc_params_validation(
1706 self, enc_params: Optional[jnp.ndarray], stash: bool = True
1707 ) -> jnp.ndarray:
1708 """
1709 Validate and normalize encoding parameters.
1711 Ensures encoding parameters have the correct shape for the model's
1712 input feature dimensions.
1714 Args:
1715 enc_params (Optional[jnp.ndarray]): Encoding parameters to validate.
1716 If None, returns the model's current encoding parameters.
1717 stash (bool): Whether to store the validated parameters on the
1718 model. See :meth:`_params_validation`. Defaults to True.
1720 Returns:
1721 jnp.ndarray: Validated encoding parameters with shape
1722 (n_qubits, n_input_feat).
1724 Raises:
1725 ValueError: If enc_params shape is incompatible with n_input_feat > 1.
1726 """
1727 if enc_params is None:
1728 enc_params = self.enc_params
1729 else:
1730 # See note in _params_validation: never stash JAX tracers on
1731 # ``self``.
1732 if stash and not isinstance(enc_params, jax.core.Tracer):
1733 if self.trainable_frequencies:
1734 self.enc_params = enc_params
1735 else:
1736 self.enc_params = jnp.array(enc_params)
1737 elif stash:
1738 log.debug(
1739 "`enc_params` is a JAX tracer; `self.enc_params` is left "
1740 "at its previous value."
1741 )
1743 if len(enc_params.shape) == 1 and self.n_input_feat == 1:
1744 enc_params = enc_params.reshape(-1, 1)
1745 elif len(enc_params.shape) == 1 and self.n_input_feat > 1:
1746 raise ValueError(
1747 f"Input dimension {self.n_input_feat} >1 but \
1748 `enc_params` has shape {enc_params.shape}"
1749 )
1751 return enc_params
1753 def _inputs_validation(
1754 self, inputs: Union[None, List, float, int, jnp.ndarray]
1755 ) -> jnp.ndarray:
1756 """
1757 Validate and normalize input data.
1759 Converts various input formats to a standardized 2D array shape
1760 suitable for batch processing in the quantum circuit.
1762 Args:
1763 inputs (Union[None, List, float, int, jnp.ndarray]): Input data in
1764 various formats:
1765 - None: Returns zeros with shape (1, n_input_feat)
1766 - float/int: Single scalar value
1767 - List: List of values or batched inputs
1768 - jnp.ndarray: NumPy/JAX array
1770 Returns:
1771 jnp.ndarray: Validated inputs with shape (batch_size, n_input_feat).
1773 Raises:
1774 ValueError: If input shape is incompatible with expected n_input_feat.
1776 Warns:
1777 UserWarning: If input is replicated to match n_input_feat.
1778 """
1779 if isinstance(inputs, List):
1780 inputs = jnp.array(np.stack(inputs))
1781 elif isinstance(inputs, float) or isinstance(inputs, int):
1782 inputs = jnp.array([inputs])
1783 elif inputs is None:
1784 inputs = jnp.array([[0] * self.n_input_feat])
1786 if len(inputs.shape) <= 1:
1787 if self.n_input_feat == 1:
1788 # add a batch dimension
1789 inputs = inputs.reshape(-1, 1)
1790 else:
1791 if inputs.shape[0] == self.n_input_feat:
1792 inputs = inputs.reshape(1, -1)
1793 else:
1794 inputs = inputs.reshape(-1, 1)
1795 inputs = inputs.repeat(self.n_input_feat, axis=1)
1796 warnings.warn(
1797 f"Expected {self.n_input_feat} inputs, but {inputs.shape[0]} "
1798 "was provided, replicating input for all input features.",
1799 UserWarning,
1800 )
1801 else:
1802 if inputs.shape[1] != self.n_input_feat:
1803 raise ValueError(
1804 f"Wrong number of inputs provided. Expected {self.n_input_feat} "
1805 f"inputs, but input has shape {inputs.shape}."
1806 )
1808 return inputs
1810 def _postprocess_res(self, result: Union[List, jnp.ndarray]) -> jnp.ndarray:
1811 """
1812 Post-process circuit execution results for uniform shape.
1814 Converts list outputs (from multiple measurements) to stacked arrays
1815 and reorders axes for consistent batch dimension placement.
1817 Args:
1818 result (Union[List, jnp.ndarray]): Raw circuit output, either a
1819 list of measurement results or a single array.
1821 Returns:
1822 jnp.ndarray: Uniformly shaped result array with batch dimension first.
1823 """
1824 if isinstance(result, list):
1825 # we use moveaxis here because in case of parity measure,
1826 # there is another dimension appended to the end and
1827 # simply transposing would result in a wrong shape
1828 result = jnp.stack(result)
1829 if len(result.shape) > 1:
1830 result = jnp.moveaxis(result, 0, 1)
1831 return result
1833 def _assimilate_batch(
1834 self,
1835 inputs: jnp.ndarray,
1836 params: jnp.ndarray,
1837 pulse_params: jnp.ndarray,
1838 enc_pulse_params: jnp.ndarray,
1839 ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, Tuple[int, ...]]:
1840 """
1841 Align batch dimensions across inputs, parameters, pulse parameters and
1842 encoding pulse parameters.
1844 Broadcasts and reshapes arrays to have compatible batch dimensions
1845 for vectorized circuit execution.
1847 The batch layout is ``[B_I, B_P, B_R, B_E, <payload>]`` where each array
1848 "owns" one batch axis and is replicated across the others (subject to
1849 the ``repeat_batch_axis`` mask) before being flattened to ``B``.
1851 Args:
1852 inputs (jnp.ndarray): Input data of shape (B_I, n_input_feat).
1853 params (jnp.ndarray): Parameters of shape (B_P, n_layers, n_params).
1854 pulse_params (jnp.ndarray): Pulse params of shape (B_R, n_layers, n_pulse).
1855 enc_pulse_params (jnp.ndarray): Encoding pulse params of shape
1856 (B_E, n_layers, n_qubits, n_enc_pulse).
1858 Returns:
1859 Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray,
1860 Tuple[int, ...]]: The four arrays, each reshaped to leading
1861 dimension B = B_I * B_P * B_R * B_E (subject to
1862 repeat_batch_axis), followed by the batch shape
1863 (B_I, B_P, B_R, B_E).
1865 Note:
1866 The effective batch shape depends on repeat_batch_axis configuration.
1867 This is the only method that derives the batch shape.
1868 """
1869 B_I = inputs.shape[0]
1870 # we check for the product because there is a chance that
1871 # there are no params. In this case we want B_P to be 1
1872 B_P = 1 if 0 in params.shape else params.shape[0]
1873 B_R = pulse_params.shape[0]
1874 B_E = enc_pulse_params.shape[0]
1876 # THIS is the only place where we derive the batch shape
1877 batch_shape = (B_I, B_P, B_R, B_E)
1878 B = np.prod(self._eff_batch_shape_of(batch_shape))
1880 # [B_I, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
1881 if B_I > 1 and self.repeat_batch_axis[0]:
1882 inputs = inputs[:, None, None, None, ...]
1883 if self.repeat_batch_axis[1]:
1884 inputs = jnp.repeat(inputs, B_P, axis=1)
1885 if self.repeat_batch_axis[2]:
1886 inputs = jnp.repeat(inputs, B_R, axis=2)
1887 if self.repeat_batch_axis[3]:
1888 inputs = jnp.repeat(inputs, B_E, axis=3)
1889 inputs = inputs.reshape(B, *inputs.shape[4:])
1891 # [B_P, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
1892 if B_P > 1 and self.repeat_batch_axis[1]:
1893 params = params[None, :, None, None, ...] # [1, B_P, 1, 1, ...]
1894 if self.repeat_batch_axis[0]:
1895 params = jnp.repeat(params, B_I, axis=0)
1896 if self.repeat_batch_axis[2]:
1897 params = jnp.repeat(params, B_R, axis=2)
1898 if self.repeat_batch_axis[3]:
1899 params = jnp.repeat(params, B_E, axis=3)
1900 params = params.reshape(B, *params.shape[4:])
1902 # [B_R, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
1903 if B_R > 1 and self.repeat_batch_axis[2]:
1904 pulse_params = pulse_params[None, None, :, None, ...] # [1, 1, B_R, 1, ...]
1905 if self.repeat_batch_axis[0]:
1906 pulse_params = jnp.repeat(pulse_params, B_I, axis=0)
1907 if self.repeat_batch_axis[1]:
1908 pulse_params = jnp.repeat(pulse_params, B_P, axis=1)
1909 if self.repeat_batch_axis[3]:
1910 pulse_params = jnp.repeat(pulse_params, B_E, axis=3)
1911 pulse_params = pulse_params.reshape(B, *pulse_params.shape[4:])
1913 # [B_E, ...] -> [B_I, B_P, B_R, B_E, ...] -> [B, ...]
1914 if B_E > 1 and self.repeat_batch_axis[3]:
1915 enc_pulse_params = enc_pulse_params[
1916 None, None, None, ...
1917 ] # [1,1,1,B_E,...]
1918 if self.repeat_batch_axis[0]:
1919 enc_pulse_params = jnp.repeat(enc_pulse_params, B_I, axis=0)
1920 if self.repeat_batch_axis[1]:
1921 enc_pulse_params = jnp.repeat(enc_pulse_params, B_P, axis=1)
1922 if self.repeat_batch_axis[2]:
1923 enc_pulse_params = jnp.repeat(enc_pulse_params, B_R, axis=2)
1924 enc_pulse_params = enc_pulse_params.reshape(B, *enc_pulse_params.shape[4:])
1926 return inputs, params, pulse_params, enc_pulse_params, batch_shape
1928 def _requires_density(self) -> bool:
1929 """
1930 Check if density matrix simulation is required.
1932 Determines whether the circuit must be executed with the mixed-state
1933 simulator based on execution type and noise configuration.
1935 Returns:
1936 bool: True if density matrix simulation is required, False otherwise.
1937 Returns True if:
1938 - execution_type is "density", or
1939 - Any non-coherent noise channel has non-zero probability
1940 """
1941 if self.execution_type == "density":
1942 return True
1944 if self.noise_params is None:
1945 return False
1947 coherent_noise = {"GateError"}
1948 for k, v in self.noise_params.items():
1949 if k in coherent_noise:
1950 continue
1951 if v is not None and v > 0:
1952 return True
1953 return False
1955 def _is_stochastic(self) -> bool:
1956 """
1957 Check if execution draws random numbers at runtime.
1959 Only coherent gate errors and shot sampling are stochastic; the Kraus
1960 channels are deterministic maps on the density matrix.
1962 Returns:
1963 bool: True if the result depends on the random key.
1964 """
1965 gate_error = (self.noise_params or {}).get("GateError") or 0
1966 return gate_error > 0 or self.shots is not None
1968 @staticmethod
1969 def _args_are_traced(*args: Any) -> bool:
1970 """
1971 Check if any argument is a JAX tracer.
1973 Args:
1974 *args (Any): Values to inspect, may be pytrees.
1976 Returns:
1977 bool: True if the call runs inside a JAX transform.
1978 """
1979 return any(
1980 isinstance(x, jax.core.Tracer) for x in jax.tree_util.tree_leaves(args)
1981 )
1983 @staticmethod
1984 def _observable_id(obs: op.Operation) -> Any:
1985 """
1986 Get a stable identity for an observable.
1988 Uses the Pauli label where available and otherwise a hash of the
1989 matrix, memoized on the instance because reading the bytes copies the
1990 full $2^n \\times 2^n$ array. Mutating a matrix in place is not
1991 detected.
1993 Args:
1994 obs (op.Operation): Observable to identify.
1996 Returns:
1997 Any: Hashable identity of the observable.
1998 """
1999 label = getattr(obs, "_pauli_label", None)
2000 if label is not None:
2001 return label
2002 if getattr(obs, "_fingerprint_hash", None) is None:
2003 obs._fingerprint_hash = hash(np.asarray(obs.matrix).tobytes())
2004 return obs._fingerprint_hash
2006 def _structural_fingerprint(self) -> Tuple:
2007 """
2008 Summarize the circuit structure for the execution plan cache.
2010 Covers everything that :meth:`_variational` and :meth:`_iec` read from
2011 the model while recording the tape and that can change after
2012 initialization without changing the shapes of the execution arguments.
2013 Without it, a batched call would silently reuse a plan that was
2014 compiled for the previous structure.
2016 Attributes that are fixed at initialization (the encoding, the state
2017 preparation, the number of qubits and layers) are omitted, as replacing
2018 them afterwards is not supported.
2020 Returns:
2021 Tuple: Hashable structure summary, passed to
2022 :meth:`~jaqsi.Script.execute`.
2023 """
2024 if self._observables is None:
2025 obs_fingerprint = None
2026 else:
2027 obs_fingerprint = tuple(
2028 (o.name, tuple(o.wires), self._observable_id(o))
2029 for o in self._observables
2030 )
2032 return (
2033 self._data_reupload.shape,
2034 # covers the derived degree, frequencies and has_dru as well
2035 self._data_reupload.tobytes(),
2036 make_hashable(self._measured_wires),
2037 obs_fingerprint,
2038 # hashed by identity, which also covers a replaced ansatz callable
2039 self.pqc,
2040 )
2042 def __call__(
2043 self,
2044 params: Optional[jnp.ndarray] = None,
2045 inputs: Optional[jnp.ndarray] = None,
2046 pulse_params: Optional[jnp.ndarray] = None,
2047 enc_params: Optional[jnp.ndarray] = None,
2048 data_reupload: Union[
2049 bool, List[List[bool]], List[List[List[bool]]], np.ndarray
2050 ] = None,
2051 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
2052 execution_type: Optional[str] = None,
2053 force_mean: bool = False,
2054 gate_mode: Optional[str] = None,
2055 enc_pulse_params: Optional[jnp.ndarray] = None,
2056 random_key: Optional[random.PRNGKey] = None,
2057 keepdims: bool = False,
2058 ) -> jnp.ndarray:
2059 """
2060 Execute the quantum circuit (callable interface).
2062 Provides a convenient callable interface for circuit execution,
2063 delegating to the _forward method.
2065 This method writes the arguments it receives onto the model, so it
2066 cannot be wrapped in an outer ``jax.jit`` or ``jax.vmap``. Use
2067 :meth:`apply` for that.
2069 Args:
2070 params (Optional[jnp.ndarray]): Variational parameters of shape
2071 (n_layers, n_params_per_layer) or (batch, n_layers, n_params_per_layer).
2072 If None, uses model's internal parameters.
2073 inputs (Optional[jnp.ndarray]): Input data of shape
2074 (batch_size, n_input_feat). If None, uses zero inputs.
2075 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
2076 the ansatz and state-preparation gates. Passing them runs those
2077 gates at pulse level. If None, they stay unitary.
2078 enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
2079 (n_qubits, n_input_feat). If None, uses model's encoding parameters.
2080 data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
2081 np.ndarray]):
2082 Data reupload configuration. If None, uses previously set reupload
2083 configuration.
2084 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
2085 Noise configuration. If None, uses previously set noise parameters.
2086 execution_type (Optional[str]): Measurement type: "expval", "density",
2087 "probs", or "state". If None, uses current execution_type setting.
2088 force_mean (bool): If True, averages results over measurement qubits.
2089 Defaults to False.
2090 gate_mode (Optional[str]): Deprecated. If None (default), the gate
2091 execution backend is inferred from the provided pulse
2092 parameters: ``pulse_params`` runs the ansatz and state
2093 preparation at pulse level, ``enc_pulse_params`` the encoding
2094 gates, both together everything. Passing "unitary",
2095 "ansatz_pulse", "enc_pulse" or "all_pulse" explicitly still
2096 works but emits a DeprecationWarning.
2097 enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
2098 for the encoding gates. Passing them runs the encoding gates at
2099 pulse level. If None, they stay unitary.
2100 random_key (Optional[random.PRNGKey]): JAX random key for stochastic
2101 execution (``GateError`` noise and shot sampling). If provided,
2102 the caller owns key advancement and the model's internal
2103 ``random_key`` is left untouched - this is the jit-safe way to
2104 get fresh randomness per call, since the internal key cannot be
2105 advanced from inside a trace. Use :meth:`next_key` to obtain
2106 one. If None, the internal key is used and advanced (eager
2107 calls only).
2108 keepdims (bool): If True, the full
2109 (B_I, B_P, B_R, B_E, O) shape is returned. If False (default),
2110 all singleton axes are squeezed out.
2112 Returns:
2113 jnp.ndarray: Circuit output with shape depending on execution_type:
2114 - "expval": (n_measured_wiress,) or scalar
2115 - "density": (2^n_output, 2^n_output)
2116 - "probs": (2^n_output,) or (n_pairs, 2^pair_size)
2117 - "state": (2^n_qubits,)
2119 Note:
2120 An eager call stores ``params``, ``pulse_params`` and ``enc_params``
2121 on the model, but a traced call (``jit``, ``grad``, ``vmap``) does
2122 not: JAX tracers must not outlive their transform, so the model
2123 state keeps its previous value. Two consequences:
2125 - Anything reading model state after a traced call - ``draw``,
2126 :class:`~qml_essentials.entanglement.Entanglement`,
2127 :class:`~qml_essentials.expressibility.Expressibility`, or a
2128 later call that omits ``params`` - sees the *old* parameters.
2129 Assign ``model.params = params`` yourself if the state should
2130 follow a traced optimization step.
2131 - Omitting ``params`` in a second call inside the same trace falls
2132 back to that stale state, so the result does not depend on the
2133 traced parameters (its gradient is zero). Pass ``params``
2134 explicitly on every call inside a trace.
2136 The skipped writes are reported at debug log level.
2137 """
2138 # Call forward method which handles the actual caching etc.
2139 return self._forward(
2140 params=params,
2141 inputs=inputs,
2142 pulse_params=pulse_params,
2143 enc_params=enc_params,
2144 data_reupload=data_reupload,
2145 noise_params=noise_params,
2146 execution_type=execution_type,
2147 force_mean=force_mean,
2148 gate_mode=gate_mode,
2149 enc_pulse_params=enc_pulse_params,
2150 random_key=random_key,
2151 keepdims=keepdims,
2152 )
2154 def apply(
2155 self,
2156 params: Optional[jnp.ndarray] = None,
2157 inputs: Optional[jnp.ndarray] = None,
2158 pulse_params: Optional[jnp.ndarray] = None,
2159 enc_params: Optional[jnp.ndarray] = None,
2160 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
2161 execution_type: Optional[str] = None,
2162 force_mean: bool = False,
2163 gate_mode: Optional[str] = None,
2164 enc_pulse_params: Optional[jnp.ndarray] = None,
2165 key: Optional[random.PRNGKey] = None,
2166 ) -> jnp.ndarray:
2167 """
2168 Execute the quantum circuit without modifying the model.
2170 Functional counterpart of :meth:`__call__`. No model state is written,
2171 so the call can be wrapped in an outer ``jax.jit``, ``jax.vmap`` or a
2172 whole jitted training step. The output always keeps the full
2173 (B_I, B_P, B_R, B_E, O) shape, so its rank does not depend on the batch
2174 sizes; call ``.squeeze()`` for the shape :meth:`__call__` returns.
2176 Arguments left as None fall back to the current model state, which an
2177 outer ``jax.jit`` bakes in at trace time. Anything that varies between
2178 calls, such as the parameters during training or the key for shots,
2179 has to be passed explicitly.
2181 Unlike :meth:`__call__` this method takes no ``data_reupload``
2182 argument, as that reconfigures the circuit; set
2183 :attr:`data_reupload` on the model beforehand instead.
2185 Args:
2186 params (Optional[jnp.ndarray]): Variational parameters of shape
2187 (n_layers, n_params_per_layer) or
2188 (batch, n_layers, n_params_per_layer).
2189 If None, uses model's internal parameters.
2190 inputs (Optional[jnp.ndarray]): Input data of shape
2191 (batch_size, n_input_feat). If None, uses zero inputs.
2192 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
2193 the ansatz and state-preparation gates. Passing them runs those
2194 gates at pulse level. If None, they stay unitary.
2195 enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
2196 (n_qubits, n_input_feat). If None, uses model's encoding parameters.
2197 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
2198 Noise configuration. If None, uses the model's noise parameters.
2199 execution_type (Optional[str]): Measurement type: "expval", "density",
2200 "probs", or "state". If None, uses current execution_type setting.
2201 force_mean (bool): If True, averages results over measurement qubits.
2202 Defaults to False.
2203 gate_mode (Optional[str]): Deprecated. If None (default), the mode
2204 is inferred from the provided pulse parameters. See
2205 :meth:`__call__`.
2206 enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
2207 for the encoding gates. Passing them runs the encoding gates at
2208 pulse level. If None, they stay unitary.
2209 key (Optional[random.PRNGKey]): JAX random key for shots and
2210 stochastic noise. If None, the model's random key is used
2211 without advancing it.
2213 Returns:
2214 jnp.ndarray: Circuit output of shape (B_I, B_P, B_R, B_E, O), where
2215 O is the per-sample output shape of the execution type and may
2216 span more than one axis (e.g. "density" and "probs").
2218 Raises:
2219 ValueError: If the encoding gates would run at pulse level but the
2220 encoding has no pulse parametrization, or if shots are set for
2221 a density measurement.
2222 """
2223 gate_mode = self._resolve_gate_mode(gate_mode, pulse_params, enc_pulse_params)
2225 if execution_type is None:
2226 execution_type = self.execution_type
2227 if execution_type == "density" and self.shots is not None:
2228 raise ValueError("Setting execution_type to density with shots not None.")
2229 result_shape = self._compute_result_shape(execution_type)
2231 if noise_params is None:
2232 noise_params = self.noise_params
2233 else:
2234 noise_params = self._normalize_noise_params(noise_params)
2236 enc_pulse_params = self._enc_pulse_params_validation(
2237 enc_pulse_params, stash=False
2238 )
2239 params = self._params_validation(params, stash=False)
2240 pulse_params = self._pulse_params_validation(pulse_params, stash=False)
2241 inputs = self._inputs_validation(inputs)
2242 enc_params = self._enc_params_validation(enc_params, stash=False)
2244 inputs, params, pulse_params, enc_pulse_params, batch_shape = (
2245 self._assimilate_batch(
2246 inputs,
2247 params,
2248 pulse_params,
2249 enc_pulse_params,
2250 )
2251 )
2253 # derive a sub key as in _forward, but without advancing the model's key
2254 _, sub_key = safe_random_split(key if key is not None else self.random_key)
2256 return self._execute_forward(
2257 params=params,
2258 inputs=inputs,
2259 pulse_params=pulse_params,
2260 enc_params=enc_params,
2261 batch_shape=batch_shape,
2262 execution_type=execution_type,
2263 result_shape=result_shape,
2264 noise_params=noise_params,
2265 gate_mode=gate_mode,
2266 force_mean=force_mean,
2267 key=sub_key,
2268 enc_pulse_params=enc_pulse_params,
2269 keepdims=True,
2270 )
2272 def _forward(
2273 self,
2274 params: Optional[jnp.ndarray] = None,
2275 inputs: Optional[jnp.ndarray] = None,
2276 pulse_params: Optional[jnp.ndarray] = None,
2277 enc_params: Optional[jnp.ndarray] = None,
2278 data_reupload: Union[
2279 bool, List[List[bool]], List[List[List[bool]]], np.ndarray
2280 ] = None,
2281 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]] = None,
2282 execution_type: Optional[str] = None,
2283 force_mean: bool = False,
2284 gate_mode: Optional[str] = None,
2285 enc_pulse_params: Optional[jnp.ndarray] = None,
2286 random_key: Optional[random.PRNGKey] = None,
2287 keepdims: bool = False,
2288 ) -> jnp.ndarray:
2289 """
2290 Execute the quantum circuit forward pass.
2292 Internal implementation of the forward pass that handles parameter
2293 validation, batch alignment, and circuit execution routing.
2295 Args:
2296 params (Optional[jnp.ndarray]): Variational parameters of shape
2297 (n_layers, n_params_per_layer) or
2298 (batch, n_layers, n_params_per_layer).
2299 If None, uses model's internal parameters.
2300 inputs (Optional[jnp.ndarray]): Input data of shape
2301 (batch_size, n_input_feat).
2302 If None, uses zero inputs.
2303 pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers for
2304 pulse-mode gate execution.
2305 enc_params (Optional[jnp.ndarray]): Encoding parameters of shape
2306 (n_qubits, n_input_feat). If None, uses model's encoding parameters.
2307 data_reupload (Union[bool, List[List[bool]], List[List[List[bool]]],
2308 np.ndarray]):
2309 Data reupload configuration. If None, uses previously set reupload
2310 configuration.
2311 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
2312 Noise configuration. If None, uses previously set noise parameters.
2313 execution_type (Optional[str]): Measurement type: "expval", "density",
2314 "probs", or "state". If None, uses current execution_type setting.
2315 force_mean (bool): If True, averages results over measurement qubits.
2316 Defaults to False.
2317 gate_mode (Optional[str]): Deprecated. If None (default), the mode
2318 is inferred from the provided pulse parameters. Passing
2319 "unitary", "ansatz_pulse", "enc_pulse" or "all_pulse"
2320 explicitly emits a DeprecationWarning.
2321 enc_pulse_params (Optional[jnp.ndarray]): Pulse parameter scalers
2322 for the encoding gates. Passing them runs the encoding gates at
2323 pulse level. If None, they stay unitary.
2324 random_key (Optional[random.PRNGKey]): JAX random key for stochastic
2325 execution. If provided, it is used instead of (and does not
2326 modify) the model's internal ``random_key``. See
2327 :meth:`__call__` for details.
2328 keepdims (bool): If True, the full
2329 (B_I, B_P, B_R, B_E, O) shape is returned. If False (default),
2330 all singleton axes are squeezed out.
2332 Returns:
2333 jnp.ndarray: Circuit output with shape depending on execution_type:
2334 - "expval": (n_measured_wiress,) or scalar
2335 - "density": (2^n_output, 2^n_output)
2336 - "probs": (2^n_output,) or (n_pairs, 2^pair_size)
2337 - "state": (2^n_qubits,)
2339 Raises:
2340 ValueError: If the encoding gates would run at pulse level but the
2341 encoding has no pulse parametrization, or if an explicitly
2342 passed (deprecated) gate_mode is unknown or inconsistent with
2343 the provided pulse parameters.
2344 """
2345 # set the parameters as object attributes
2346 if noise_params is not None:
2347 self.noise_params = noise_params
2348 if execution_type is not None:
2349 self.execution_type = execution_type
2351 # one frame deeper than apply, as __call__ sits in between
2352 gate_mode = self._resolve_gate_mode(
2353 gate_mode, pulse_params, enc_pulse_params, stacklevel=4
2354 )
2356 # TODO: add testing
2357 if data_reupload is not None:
2358 self.data_reupload = data_reupload
2360 params = self._params_validation(params)
2361 pulse_params = self._pulse_params_validation(pulse_params)
2362 inputs = self._inputs_validation(inputs)
2363 enc_params = self._enc_params_validation(enc_params)
2364 enc_pulse_params = self._enc_pulse_params_validation(enc_pulse_params)
2366 inputs, params, pulse_params, enc_pulse_params, batch_shape = (
2367 self._assimilate_batch(
2368 inputs,
2369 params,
2370 pulse_params,
2371 enc_pulse_params,
2372 )
2373 )
2374 self._batch_shape = batch_shape
2376 # split to generate a sub_key, required for actual execution.
2377 if random_key is not None:
2378 # explicit key: purely functional, the caller advances it. This is
2379 # the only way to get fresh randomness inside a trace, because a
2380 # jitted call is traced once and then replays the trace-time key.
2381 _, sub_key = safe_random_split(random_key)
2382 else:
2383 if self._is_stochastic() and self._args_are_traced(
2384 params, inputs, pulse_params, enc_pulse_params
2385 ):
2386 warnings.warn(
2387 "Stochastic execution (`GateError` or `shots`) without an "
2388 "explicit `random_key` inside a JAX transform: the key is "
2389 "read at trace time, so a jitted function replays the same "
2390 "noise realization on every call. Pass "
2391 "`random_key=model.next_key()` from outside the transform.",
2392 UserWarning,
2393 )
2394 # Under JAX tracing (jit) the split result is a tracer; stashing it
2395 # on ``self`` leaks the tracer across calls (UnexpectedTracerError),
2396 # so only advance the key eagerly. Note that a jitted call
2397 # therefore reuses the same key on every execution - pass
2398 # ``random_key`` explicitly if that matters.
2399 new_key, sub_key = safe_random_split(self.random_key)
2400 if not isinstance(new_key, jax.core.Tracer):
2401 self.random_key = new_key
2403 return self._execute_forward(
2404 params=params,
2405 inputs=inputs,
2406 pulse_params=pulse_params,
2407 enc_params=enc_params,
2408 batch_shape=batch_shape,
2409 enc_pulse_params=enc_pulse_params,
2410 execution_type=self.execution_type,
2411 result_shape=self._result_shape,
2412 noise_params=self.noise_params,
2413 gate_mode=gate_mode,
2414 force_mean=force_mean,
2415 key=sub_key,
2416 keepdims=keepdims,
2417 )
2419 def _execute_forward(
2420 self,
2421 params: jnp.ndarray,
2422 inputs: jnp.ndarray,
2423 pulse_params: jnp.ndarray,
2424 enc_params: jnp.ndarray,
2425 enc_pulse_params: jnp.ndarray,
2426 batch_shape: Tuple[int, ...],
2427 execution_type: str,
2428 result_shape: Tuple[int, ...],
2429 noise_params: Optional[Dict[str, Union[float, Dict[str, float]]]],
2430 gate_mode: str,
2431 force_mean: bool,
2432 key: random.PRNGKey,
2433 keepdims: bool,
2434 ) -> jnp.ndarray:
2435 """
2436 Run the simulation and post-process the result.
2438 This is the shared core of :meth:`_forward` and :meth:`apply`. Every
2439 value it needs is passed in explicitly and no model state is written,
2440 so it is safe to call from within an outer JAX transform.
2442 Args:
2443 params (jnp.ndarray): Validated and batch-aligned parameters.
2444 inputs (jnp.ndarray): Validated and batch-aligned inputs.
2445 pulse_params (jnp.ndarray): Validated and batch-aligned pulse params.
2446 enc_params (jnp.ndarray): Validated encoding parameters.
2447 enc_pulse_params (jnp.ndarray): Validated and batch-aligned
2448 encoding pulse parameters.
2449 batch_shape (Tuple[int, ...]): Batch shape (B_I, B_P, B_R, B_E).
2450 execution_type (str): Measurement type: "expval", "density",
2451 "probs", or "state".
2452 result_shape (Tuple[int, ...]): Per-sample output shape.
2453 noise_params (Optional[Dict[str, Union[float, Dict[str, float]]]]):
2454 Normalized noise configuration.
2455 gate_mode (str): Gate execution backend, "unitary",
2456 "ansatz_pulse", "enc_pulse" or "all_pulse".
2457 force_mean (bool): If True, averages results over the output axis.
2458 key (random.PRNGKey): JAX random key for this execution.
2459 keepdims (bool): If True, the full (B_I, B_P, B_R, B_E, O) shape is
2460 returned. If False, all singleton axes are squeezed out.
2462 Returns:
2463 jnp.ndarray: Circuit output.
2464 """
2465 # Build measurement type & observables from execution_type / output_qubit
2466 meas_type, obs = self._build_obs(execution_type)
2468 eff_batch_shape = self._eff_batch_shape_of(batch_shape)
2470 # Jaqsi auto-routes between statevector and density-matrix simulation
2471 # based on whether noise channels appear on the tape, so a single
2472 B = np.prod(eff_batch_shape)
2474 # kwargs are broadcast (not vmapped over)
2475 exec_kwargs = dict(
2476 noise_params=noise_params,
2477 gate_mode=gate_mode,
2478 )
2480 # Build a shot key from the given key if shots are requested
2481 shot_key = None
2482 sub_key = key
2483 if self.shots is not None:
2484 # overwrite subkey and split shot_key
2485 sub_key, shot_key = safe_random_split(sub_key)
2487 if B > 1:
2488 # use random keys, derived from the subkey
2489 random_keys = safe_random_split(sub_key, num=B)
2491 in_axes = (
2492 0 if batch_shape[1] > 1 else None, # params
2493 0 if batch_shape[0] > 1 else None, # inputs
2494 0 if batch_shape[2] > 1 else None, # pulse_params
2495 0, # random_keys
2496 None, # enc_params (broadcast, not batched)
2497 0 if batch_shape[3] > 1 else None, # enc_pulse_params
2498 )
2500 result = self.script.execute(
2501 type=meas_type,
2502 obs=obs,
2503 args=(
2504 params,
2505 inputs,
2506 pulse_params,
2507 random_keys,
2508 enc_params,
2509 enc_pulse_params,
2510 ),
2511 kwargs=exec_kwargs,
2512 in_axes=in_axes,
2513 shots=self.shots,
2514 key=shot_key,
2515 fingerprint=self._structural_fingerprint(),
2516 )
2517 else:
2518 # use the subkey directly
2519 result = self.script.execute(
2520 type=meas_type,
2521 obs=obs,
2522 args=(
2523 params,
2524 inputs,
2525 pulse_params,
2526 sub_key,
2527 enc_params,
2528 enc_pulse_params,
2529 ),
2530 kwargs=exec_kwargs,
2531 shots=self.shots,
2532 key=shot_key,
2533 fingerprint=self._structural_fingerprint(),
2534 )
2536 result = self._postprocess_res(result)
2538 # --- Post-processing for partial-qubit measurements ---------------
2539 if execution_type == "density" and not self.all_qubit_measurement:
2540 result = js.partial_trace(result, self.n_qubits, self._measured_wires)
2542 if execution_type == "probs" and not self.all_qubit_measurement:
2543 if isinstance(self._measured_wires[0], (list, tuple)):
2544 # list of qubit groups - marginalize each independently
2545 result = jnp.stack(
2546 [
2547 js.marginalize_probs(result, self.n_qubits, list(group))
2548 for group in self._measured_wires
2549 ]
2550 )
2551 else:
2552 result = js.marginalize_probs(
2553 result, self.n_qubits, self._measured_wires
2554 )
2556 result = jnp.asarray(result)
2557 result = result.reshape((*eff_batch_shape, *result_shape))
2558 if not keepdims:
2559 result = result.squeeze()
2561 if (
2562 execution_type in ("expval", "probs")
2563 and force_mean
2564 and len(result.shape) > 0
2565 and result_shape[0] > 1
2566 ):
2567 result = result.mean(axis=-1, keepdims=keepdims)
2569 return result