From 81d4255d5bc226620dc584491bb5a71f603e9441 Mon Sep 17 00:00:00 2001 From: Will Date: Thu, 8 Oct 2026 11:18:03 -0400 Subject: [PATCH 1/7] first commit --- pennylane/labs/tcdq/__init__.py | 35 +-- pennylane/labs/tcdq/mmd_loss_hw.py | 2 +- pennylane/labs/tcdq/qudit_expval_functions.py | 251 ++++++++++++------ 3 files changed, 190 insertions(+), 98 deletions(-) diff --git a/pennylane/labs/tcdq/__init__.py b/pennylane/labs/tcdq/__init__.py index 7dff655947e..96fa31272fe 100644 --- a/pennylane/labs/tcdq/__init__.py +++ b/pennylane/labs/tcdq/__init__.py @@ -194,9 +194,10 @@ def loss_fn(current_params): The same workflow applies to qudit circuits (``d > 2``). Replace :class:`~CircuitConfig` with :class:`~QuditCircuitConfig` and :func:`~build_expval_func` with :func:`~build_qudit_expval_func`. -Gate vectors now have length ``n_qudits`` with entries in -:math:`\{0, \ldots, d-1\}` specifying the power of :math:`Z` on each -qudit. +Each gate is a sparse ``dict`` mapping a qudit index to the power of +:math:`Z` applied on that qudit (entries in :math:`\{1, \ldots, d-1\}`); +qudits that are not listed are left untouched, so a :math:`k`-local gate +costs :math:`O(k)` to describe regardless of ``n_qudits``. .. code-block:: python @@ -210,13 +211,13 @@ def loss_fn(current_params): # Single-qudit and nearest-neighbour two-qudit gates gates = { - 0: [[1, 0, 0, 0]], - 1: [[0, 1, 0, 0]], - 2: [[0, 0, 1, 0]], - 3: [[0, 0, 0, 1]], - 4: [[1, 1, 0, 0]], - 5: [[0, 1, 1, 0]], - 6: [[0, 0, 1, 1]], + 0: [{0: 1}], + 1: [{1: 1}], + 2: [{2: 1}], + 3: [{3: 1}], + 4: [{0: 1, 1: 1}], + 5: [{1: 1, 2: 1}], + 6: [{2: 1, 3: 1}], } # Observables: displacement operators O(l, m) with m = 0 @@ -270,13 +271,13 @@ def loss_fn(current_params): # Define single-qudit and nearest-neighbour two-qudit gates gates = { - 0: [[1, 0, 0, 0]], - 1: [[0, 1, 0, 0]], - 2: [[0, 0, 1, 0]], - 3: [[0, 0, 0, 1]], - 4: [[1, 1, 0, 0]], - 5: [[0, 1, 1, 0]], - 6: [[0, 0, 1, 1]], + 0: [{0: 1}], + 1: [{1: 1}], + 2: [{2: 1}], + 3: [{3: 1}], + 4: [{0: 1, 1: 1}], + 5: [{1: 1, 2: 1}], + 6: [{2: 1, 3: 1}], } circuit_config = QuditCircuitConfig( diff --git a/pennylane/labs/tcdq/mmd_loss_hw.py b/pennylane/labs/tcdq/mmd_loss_hw.py index 09823c2db95..1d48bf98205 100644 --- a/pennylane/labs/tcdq/mmd_loss_hw.py +++ b/pennylane/labs/tcdq/mmd_loss_hw.py @@ -393,7 +393,7 @@ def build_mmd_loss_hw( >>> circuit_config = QuditCircuitConfig( ... dims=3, ... n_qudits=2, - ... gates={0: [[1, 0]], 1: [[0, 1]]}, + ... gates={0: [{0: 1}], 1: [{1: 1}]}, ... n_samples=512, ... key=jax.random.PRNGKey(0), ... ) diff --git a/pennylane/labs/tcdq/qudit_expval_functions.py b/pennylane/labs/tcdq/qudit_expval_functions.py index 27b2083e316..17bda6ad66e 100644 --- a/pennylane/labs/tcdq/qudit_expval_functions.py +++ b/pennylane/labs/tcdq/qudit_expval_functions.py @@ -71,14 +71,18 @@ class QuditCircuitConfig: # pylint: disable=too-many-instance-attributes every qudit, or a sequence of length ``n_qudits`` giving a distinct dimension :math:`d_j` per qudit. n_qudits (int): Number of qudits in the circuit. - gates (dict[int, list[list[int]]]): Circuit structure mapping each - trainable-parameter index to a list of generator vectors. Each - generator vector has length ``n_qudits`` with integer entries in - :math:`\{0, \ldots, d_j-1\}` that specify the power of :math:`Z` on - each qudit. For example, with ``d=3`` and ``n_qudits=2``, - ``{0: [[1, 0]], 1: [[0, 1]], 2: [[1, 1]]}`` defines three gates: - :math:`Z^1` on qudit 0, :math:`Z^1` on qudit 1, and - :math:`Z^1 \otimes Z^1` on both. + gates (dict[int, list[dict[int, int]]]): Circuit structure mapping each + trainable-parameter index to a list of gates. Each gate is a sparse + generator given as a ``dict`` mapping a qudit index to the power of + :math:`Z` applied on that qudit; qudits that are not listed are + acted on trivially. Powers are reduced modulo the local dimension + :math:`d_j`, and entries that reduce to zero are dropped. For + example, with ``dims=3`` and ``n_qudits=2``, + ``{0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]}`` defines three + gates: :math:`Z^1` on qudit 0, :math:`Z^1` on qudit 1, and + :math:`Z^1 \otimes Z^1` on both. Only the active qudits of each + gate are stored, so the cost of describing a :math:`k`-local + circuit scales with :math:`k` rather than with ``n_qudits``. n_samples (int): Number of random dit-strings drawn for the estimation. key (ArrayLike): JAX PRNG key for random dit-string generation. @@ -107,9 +111,9 @@ class QuditCircuitConfig: # pylint: disable=too-many-instance-attributes >>> import jax.numpy as jnp >>> from pennylane.labs.tcdq import QuditCircuitConfig >>> config = QuditCircuitConfig( - ... d=3, + ... dims=3, ... n_qudits=4, - ... gates={0: [[1, 0, 0, 0]], 1: [[0, 1, 0, 0]], 2: [[1, 1, 0, 0]]}, + ... gates={0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]}, ... observables=( ... jnp.array([[1, 0, 0, 0], [0, 1, 0, 0]], dtype=jnp.int32), ... jnp.zeros((2, 4), dtype=jnp.int32), @@ -123,8 +127,8 @@ class QuditCircuitConfig: # pylint: disable=too-many-instance-attributes dims: int | Sequence[int] = None #: Number of qudits in the circuit. n_qudits: int = None - #: Circuit structure mapping parameter indices to generator vectors. - gates: dict[int, list[list[int]]] = None + #: Circuit structure mapping parameter indices to sparse gates ``{qudit: power}``. + gates: dict[int, list[dict[int, int]]] = None #: Number of random dit-strings drawn for the estimation. n_samples: int = None #: JAX PRNG key for random dit-string generation. @@ -162,44 +166,151 @@ def _dims_to_numpy(dims: int | Sequence[int], n_qudits: int) -> np.ndarray: return normalized_dims -def _parse_qudit_generator_dict(circuit_def: dict[int, list[list[int]]], n_qudits: int): - """Convert a qudit gate dictionary into a generator matrix and parameter map. +class SparseGateGroup(NamedTuple): + """Gates of equal weight :math:`\\omega` stored in sparse ``(support, power)`` form. - Unlike the qubit version, generator vectors are provided explicitly (not as wire - indices), so each inner list must already have length ``n_qudits`` with integer entries - in ``{0, ..., d-1}``. + Only the active qudits of each gate are kept, so the storage cost is + ``O(n_gates * omega)`` instead of ``O(n_gates * n_qudits)``. Args: - circuit_def (dict[int, list[list[int]]]): Maps parameter indices to lists of - generator vectors of length ``n_qudits``. + omega (int): Number of active qudits shared by every gate in the group. + supports (np.ndarray): Active qudit indices, shape ``(n_gates, omega)``, + sorted in increasing order along the last axis. + powers (np.ndarray): Power of :math:`Z` at each support position, shape + ``(n_gates, omega)``, with entries in ``{1, ..., d_j - 1}``. + param_indices (jnp.ndarray): Index into ``gates_params`` for each gate, + shape ``(n_gates,)``. + """ + + #: Number of active qudits per gate in this group. + omega: int + #: Active qudit indices, shape ``(n_gates, omega)``. + supports: np.ndarray + #: Power of :math:`Z` at each support position, shape ``(n_gates, omega)``. + powers: np.ndarray + #: Parameter index of each gate, shape ``(n_gates,)``. + param_indices: jnp.ndarray + + +def _normalize_sparse_gate( + gate: dict[int, int], n_qudits: int, dims: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Validate one sparse gate and return its sorted ``(support, powers)`` arrays. + + Qudit indices may be negative (Python-style, counted from the end). Powers + are reduced modulo the local dimension of their qudit and zero powers are + dropped, so the returned support contains only genuinely active qudits. + + Raises: + TypeError: If ``gate`` is not a ``dict`` (e.g. a dense length-``n_qudits`` vector). + IndexError: If a qudit index is outside ``[-n_qudits, n_qudits)``. + """ + if not isinstance(gate, dict): + raise TypeError( + "Each qudit gate must be a dict mapping qudit index to Z power, e.g. {0: 1, 2: 2}; " + f"got {type(gate).__name__}. Dense length-n_qudits generator vectors are not accepted." + ) + + support: list[int] = [] + powers: list[int] = [] + for qudit, power in gate.items(): + q = int(qudit) + if q < -n_qudits or q >= n_qudits: + raise IndexError(f"Qudit index {q} out of range for a {n_qudits}-qudit circuit.") + q = q % n_qudits + p = int(power) % int(dims[q]) + if p == 0: + continue + support.append(q) + powers.append(p) + + order = np.argsort(support, kind="stable") + return np.asarray(support, dtype=int)[order], np.asarray(powers, dtype=int)[order] + + +def _parse_qudit_gate_dict( + circuit_def: dict[int, list[dict[int, int]]], n_qudits: int, dims: ArrayLike +) -> list[SparseGateGroup]: + """Convert a sparse qudit gate dictionary into weight-grouped support/power arrays. + + This is the qudit analogue of + :func:`~pennylane.labs.tcdq.expval_functions._parse_generator_dict`. It never + materialises a dense ``(n_gates, n_qudits)`` generator matrix: each gate is + stored through its active qudits only, and gates are bucketed by weight so + that the downstream :math:`2^\\omega` angle-addition expansion can be + vectorised within each bucket. + + Gates whose support is empty after reducing powers modulo ``dims`` act as + the identity and are dropped. Their parameter index is still a valid slot + in ``gates_params``; it simply receives zero gradient. + + Args: + circuit_def (dict[int, list[dict[int, int]]]): Maps parameter indices to lists + of sparse gates ``{qudit_index: z_power}``. n_qudits (int): Number of qudits. + dims (ArrayLike): Per-qudit dimensions, shape ``(n_qudits,)``. Returns: - tuple[jnp.ndarray, jnp.ndarray]: Tuple containing: - - Integer generator matrix of shape ``(n_gates, n_qudits)``. - - Integer array mapping each gate to its parameter index. + list[SparseGateGroup]: One group per distinct non-zero weight, sorted by weight. Raises: - ValueError: If any generator vector has length != ``n_qudits``. + TypeError: If a gate is not a ``dict``. + IndexError: If a qudit index is out of range. """ - flat_gates = [] - param_indices = [] + dims = np.asarray(dims, dtype=int) + buckets: dict[int, tuple[list[np.ndarray], list[np.ndarray], list[int]]] = {} for param_idx in sorted(circuit_def.keys()): for gate in circuit_def[param_idx]: - if len(gate) != n_qudits: - raise ValueError(f"Generator has length {len(gate)}, expected {n_qudits}.") - flat_gates.append(gate) - param_indices.append(param_idx) - - n_gates = len(flat_gates) - if n_gates == 0: - generators = np.zeros((0, n_qudits), dtype=int) - else: - generators = np.array(flat_gates, dtype=int) + support, powers = _normalize_sparse_gate(gate, n_qudits, dims) + omega = support.size + if omega == 0: + continue + sup_list, pow_list, pidx_list = buckets.setdefault(omega, ([], [], [])) + sup_list.append(support) + pow_list.append(powers) + pidx_list.append(int(param_idx)) + + return [ + SparseGateGroup( + omega=omega, + supports=np.stack(sup_list), + powers=np.stack(pow_list), + param_indices=jnp.array(pidx_list, dtype=int), + ) + for omega, (sup_list, pow_list, pidx_list) in sorted(buckets.items()) + ] + + +def _gates_to_dense_generators( + circuit_def: dict[int, list[dict[int, int]]], n_qudits: int, dims: ArrayLike +) -> tuple[np.ndarray, np.ndarray]: + """Expand sparse gates into a dense generator matrix (for brute-force references). + + This is intentionally *not* used by the estimator; it exists so that + exhaustive reference implementations and tests can consume the same + ``gates`` dictionary. Gates are flattened in the same order as + :func:`_parse_qudit_gate_dict` visits them (sorted parameter index, then + list order) and identity gates are kept so that the returned + ``param_map`` has one entry per gate in ``circuit_def``. + + Returns: + tuple[np.ndarray, np.ndarray]: ``(generators, param_map)`` with shapes + ``(n_gates, n_qudits)`` and ``(n_gates,)``. + """ + dims = np.asarray(dims, dtype=int) + rows: list[np.ndarray] = [] + param_map: list[int] = [] + for param_idx in sorted(circuit_def.keys()): + for gate in circuit_def[param_idx]: + support, powers = _normalize_sparse_gate(gate, n_qudits, dims) + row = np.zeros((n_qudits,), dtype=int) + row[support] = powers + rows.append(row) + param_map.append(int(param_idx)) - param_map = jnp.array(param_indices, dtype=int) - return jnp.array(generators), param_map + generators = np.stack(rows) if rows else np.zeros((0, n_qudits), dtype=int) + return generators, np.asarray(param_map, dtype=int) def _compute_qudit_samples( @@ -315,30 +426,32 @@ def _expand_angle_addition( def _build_weight_group( - generators_w: np.ndarray, - param_indices: jnp.ndarray, + group: SparseGateGroup, samples: jnp.ndarray, l_vecs: jnp.ndarray, dims: np.ndarray, ) -> WeightGroupData: - """Precompute the factor matrices for a group of gates with the same weight.""" - n_gates = len(generators_w) + """Precompute the factor matrices for a group of gates with the same weight. + + The supports and powers come straight from the sparse gate description, so + no dense generator matrix is scanned here. + """ + n_gates, omega = group.supports.shape num_samples = samples.shape[0] n_obs = l_vecs.shape[0] - omega = int(np.count_nonzero(generators_w[0])) - supports = np.array([np.where(g != 0)[0] for g in generators_w]) # (n_gates, omega) - gate_vals = np.array([g[s] for g, s in zip(generators_w, supports)]) # (n_gates, omega) - d_at_support = np.asarray(dims)[supports] # (n_gates, omega) + d_at_support = np.asarray(dims)[group.supports] # (n_gates, omega) - z_at_support = _gather_support_values(samples, supports, num_samples, n_gates, omega) - l_at_support = _gather_support_values(l_vecs, supports, n_obs, n_gates, omega) + z_at_support = _gather_support_values(samples, group.supports, num_samples, n_gates, omega) + l_at_support = _gather_support_values(l_vecs, group.supports, n_obs, n_gates, omega) state_cos, state_sin, obs_cos, obs_sin = _compute_trigonometric_building_blocks( - gate_vals, z_at_support, l_at_support, d_at_support + group.powers, z_at_support, l_at_support, d_at_support ) samples_matrices, obs_matrices = _expand_angle_addition(state_cos, state_sin, obs_cos, obs_sin) return WeightGroupData( - param_indices=param_indices, samples_matrices=samples_matrices, obs_matrices=obs_matrices + param_indices=group.param_indices, + samples_matrices=samples_matrices, + obs_matrices=obs_matrices, ) @@ -368,33 +481,17 @@ def _obs_phase_matrix( ) -# pylint: disable=too-many-arguments def _build_all_weight_groups( - gen_np: np.ndarray, - pm_np: np.ndarray, - gate_weights: np.ndarray, + gate_groups: list[SparseGateGroup], samples: jnp.ndarray, l_vecs: jnp.ndarray, dims: np.ndarray, ) -> list[WeightGroupData]: - """Build :class:`WeightGroupData` for each non-zero gate weight.""" - weight_data: list[WeightGroupData] = [] - for omega in sorted(set(gate_weights)): - if omega == 0: - continue - gate_indices = np.where(gate_weights == omega)[0] - weight_data.append( - _build_weight_group( - generators_w=gen_np[gate_indices], - param_indices=jnp.array(pm_np[gate_indices]), - samples=samples, - l_vecs=l_vecs, - dims=dims, - ) - ) - return weight_data + """Build :class:`WeightGroupData` for each sparse gate group.""" + return [_build_weight_group(group, samples, l_vecs, dims) for group in gate_groups] +# pylint: disable=too-many-arguments def _accumulate_phase_diffs( gates_params: ArrayLike, weight_data: list[WeightGroupData], @@ -525,9 +622,9 @@ def build_qudit_expval_func( # pylint: disable=too-many-statements >>> import jax.numpy as jnp >>> from pennylane.labs.tcdq import QuditCircuitConfig, build_qudit_expval_func >>> config = QuditCircuitConfig( - ... d=3, + ... dims=3, ... n_qudits=2, - ... gates={0: [[1, 0]], 1: [[0, 1]]}, + ... gates={0: [{0: 1}], 1: [{1: 1}]}, ... n_samples=512, ... key=jax.random.PRNGKey(0), ... observables=( @@ -545,10 +642,9 @@ def build_qudit_expval_func( # pylint: disable=too-many-statements `Spectral Born machines: classically trainable quantum generative models for discrete data `_. """ - generators, param_map = _parse_qudit_generator_dict(config.gates, config.n_qudits) - n = config.n_qudits dims = _dims_to_numpy(config.dims, n) + gate_groups = _parse_qudit_gate_dict(config.gates, n, dims) default_samples = _compute_qudit_samples(config.key, config.n_samples, n, dims) vmapped_phase_func = None @@ -565,9 +661,6 @@ def compute_phase_diff(p_params, sample, l_vec): in_axes=(None, None, 0), ) - gen_np, pm_np = np.array(generators), np.array(param_map) - gate_weights = np.sum(gen_np != 0, axis=1) - if config.observables is not None: l_vecs = jnp.array(config.observables[0], dtype=jnp.int32) m_vecs = jnp.array(config.observables[1], dtype=jnp.int32) @@ -579,9 +672,7 @@ def compute_phase_diff(p_params, sample, l_vec): n_obs=n_obs, l_f=l_f, m_f=m_f, - weight_data=_build_all_weight_groups( - gen_np, pm_np, gate_weights, default_samples, l_vecs, dims - ), + weight_data=_build_all_weight_groups(gate_groups, default_samples, l_vecs, dims), obs_phase_matrix=_obs_phase_matrix(default_samples, m_f, l_f, dims), ) else: @@ -651,7 +742,7 @@ def qudit_expval_batched( w_data = defaults.weight_data else: obs_pm = _obs_phase_matrix(samples, m_f, l_f, dims) - w_data = _build_all_weight_groups(gen_np, pm_np, gate_weights, samples, l_vecs, dims) + w_data = _build_all_weight_groups(gate_groups, samples, l_vecs, dims) accumulated_phase_diffs = _accumulate_phase_diffs( gates_params, w_data, n_obs, _n, vmapped_phase_func, phase_fn_params, samples, l_vecs From 89f818a528ee8d19479c1045a15fdd14b438e397 Mon Sep 17 00:00:00 2001 From: Will Date: Thu, 8 Oct 2026 11:36:55 -0400 Subject: [PATCH 2/7] update tests --- pennylane/labs/tests/tcdq/test_mmd_loss_hw.py | 88 +++--- .../labs/tests/tcdq/test_qudit_expval.py | 257 ++++++++++++------ 2 files changed, 211 insertions(+), 134 deletions(-) diff --git a/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py b/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py index 77e0555737c..e836141b2ed 100644 --- a/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py +++ b/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py @@ -31,6 +31,7 @@ from pennylane.labs.tcdq.qudit_expval_functions import ( QuditCircuitConfig, _dims_to_numpy, + _gates_to_dense_generators, build_qudit_expval_func, ) @@ -72,15 +73,10 @@ def _qudit_phi_g_z(gen, z, d): return val -def _flatten_gates(gates, n_qudits): - """Flatten the gate dictionary into generators plus parameter indices.""" - flat_gens, flat_pidx = [], [] - for pidx in sorted(gates.keys()): - for g in gates[pidx]: - assert len(g) == n_qudits - flat_gens.append(np.array(g, dtype=int)) - flat_pidx.append(pidx) - return flat_gens, flat_pidx +def _flatten_gates(gates, n_qudits, d): + """Flatten the sparse gate dictionary into dense generators plus parameter indices.""" + generators, flat_pidx = _gates_to_dense_generators(gates, n_qudits, _dims_to_numpy(d, n_qudits)) + return [np.array(g, dtype=int) for g in generators], [int(i) for i in flat_pidx] def _qudit_circuit_probs(gates, params, n_qudits, d, phase_fn=None, phase_params=None): @@ -91,8 +87,8 @@ def _qudit_circuit_probs(gates, params, n_qudits, d, phase_fn=None, phase_params gate layer, matching ``QuditCircuitConfig.phase_fn``. """ params = np.asarray(params, dtype=float) - flat_gens, flat_pidx = _flatten_gates(gates, n_qudits) dims = _dims_to_numpy(d, n_qudits) + flat_gens, flat_pidx = _flatten_gates(gates, n_qudits, d) omega = np.exp(2j * np.pi / dims) # per-qudit roots of unity, shape (n_qudits,) all_states = list(itertools.product(*(range(int(d_j)) for d_j in dims))) n_states = int(np.prod(dims)) @@ -123,8 +119,8 @@ def _qudit_circuit_probs(gates, params, n_qudits, d, phase_fn=None, phase_params def _qudit_expval_exact(gates, params, l_vec, n_qudits, d): """Return one exact observable moment by summing over all basis states.""" params = np.asarray(params, dtype=float) - flat_gens, flat_pidx = _flatten_gates(gates, n_qudits) dims = _dims_to_numpy(d, n_qudits) + flat_gens, flat_pidx = _flatten_gates(gates, n_qudits, d) all_states = list(itertools.product(*(range(int(d_j)) for d_j in dims))) n_states = int(np.prod(dims)) l_arr = np.array(l_vec) @@ -381,21 +377,21 @@ class TestExactQuditMMDConsistency: "gates, params, n_qudits, d, graph_type", [ ( - {0: [[1, 0]], 1: [[0, 1]], 2: [[1, 1]]}, + {0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]}, [0.3, 0.5, 0.2], 2, 3, "cycle", ), ( - {0: [[1, 0]], 1: [[0, 2]]}, + {0: [{0: 1}], 1: [{1: 2}]}, [0.4, 0.7], 2, 3, "complete", ), ( - {0: [[1, 0]], 1: [[0, 1]], 2: [[2, 1]]}, + {0: [{0: 1}], 1: [{1: 1}], 2: [{0: 2, 1: 1}]}, [0.1, 0.6, 0.9], 2, 4, @@ -403,7 +399,7 @@ class TestExactQuditMMDConsistency: ), # Non-uniform: qubit x qutrit, cycle ( - {0: [[1, 0]], 1: [[0, 2]], 2: [[1, 1]]}, + {0: [{0: 1}], 1: [{1: 2}], 2: [{0: 1, 1: 1}]}, [0.3, 0.5, 0.2], 2, [2, 3], @@ -411,7 +407,7 @@ class TestExactQuditMMDConsistency: ), # Non-uniform: qutrit x ququart, complete ( - {0: [[1, 0]], 1: [[0, 3]]}, + {0: [{0: 1}], 1: [{1: 3}]}, [0.4, 0.7], 2, [3, 4], @@ -441,7 +437,7 @@ def test_raises_target_data_too_few_samples(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -454,7 +450,7 @@ def test_raises_n_ops_zero(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -467,7 +463,7 @@ def test_raises_empty_bandwidth(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -480,7 +476,7 @@ def test_raises_wire_out_of_range(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -493,7 +489,7 @@ def test_raises_duplicate_wires(self): config = QuditCircuitConfig( dims=3, n_qudits=3, - gates={0: [[1, 0, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -506,7 +502,7 @@ def test_raises_target_data_wrong_ndim(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -519,7 +515,7 @@ def test_raises_target_data_wrong_columns(self): config = QuditCircuitConfig( dims=3, n_qudits=3, - gates={0: [[1, 0, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -533,7 +529,7 @@ def test_return_per_bandwidth_type_and_length(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -547,7 +543,7 @@ def test_multi_bandwidth_mean(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=200, key=jax.random.PRNGKey(42), ) @@ -575,7 +571,7 @@ def test_single_bandwidth_returns_scalar(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]]}, + gates={0: [{0: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -589,7 +585,7 @@ def test_deterministic_same_key(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]], 2: [[1, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]}, n_samples=100, key=jax.random.PRNGKey(99), ) @@ -606,7 +602,7 @@ def test_different_keys_give_different_results(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -627,7 +623,7 @@ def test_mmd_loss_with_custom_init_state(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 1]]}, + gates={0: [{0: 1, 1: 1}]}, n_samples=100, key=jax.random.PRNGKey(42), init_state_elems=state_elems, @@ -658,7 +654,7 @@ class TestQuditMMDLossStatistical: [ # 2-qutrit, 3 gates, cycle ( - {0: [[1, 0]], 1: [[0, 1]], 2: [[1, 1]]}, + {0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]}, [0.37, 0.95, 0.73], 2, 3, @@ -667,7 +663,7 @@ class TestQuditMMDLossStatistical: ), # 2-qutrit, 2 gates, complete ( - {0: [[1, 0]], 1: [[0, 2]]}, + {0: [{0: 1}], 1: [{1: 2}]}, [0.5, 0.3], 2, 3, @@ -676,7 +672,7 @@ class TestQuditMMDLossStatistical: ), # 2-qudit d=4, 3 gates, cycle ( - {0: [[1, 0]], 1: [[0, 1]], 2: [[2, 1]]}, + {0: [{0: 1}], 1: [{1: 1}], 2: [{0: 2, 1: 1}]}, [0.2, 0.8, 0.4], 2, 4, @@ -685,7 +681,7 @@ class TestQuditMMDLossStatistical: ), # Non-uniform: qubit x qutrit, cycle ( - {0: [[1, 0]], 1: [[0, 2]], 2: [[1, 1]]}, + {0: [{0: 1}], 1: [{1: 2}], 2: [{0: 1, 1: 1}]}, [0.37, 0.95, 0.73], 2, [2, 3], @@ -694,7 +690,7 @@ class TestQuditMMDLossStatistical: ), # Non-uniform: qutrit x ququart, complete ( - {0: [[1, 0]], 1: [[0, 3]]}, + {0: [{0: 1}], 1: [{1: 3}]}, [0.5, 0.3], 2, [3, 4], @@ -754,7 +750,7 @@ def test_wires_subset_executes(self): config = QuditCircuitConfig( dims=3, n_qudits=3, - gates={0: [[1, 0, 0]], 1: [[0, 1, 0]], 2: [[0, 0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}], 2: [{2: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -769,7 +765,7 @@ def test_sqrt_loss_positive(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=100, key=jax.random.PRNGKey(77), ) @@ -784,7 +780,7 @@ def test_key_override_provides_new_randomness(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=100, key=jax.random.PRNGKey(0), ) @@ -804,7 +800,7 @@ def test_both_graph_types_finite(self, graph_type): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=200, key=jax.random.PRNGKey(0), ) @@ -818,7 +814,7 @@ def test_gradient_flows(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=200, key=jax.random.PRNGKey(0), ) @@ -838,7 +834,7 @@ def _make_config_and_data(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]], 2: [[1, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]}, n_samples=200, key=jax.random.PRNGKey(42), ) @@ -907,7 +903,7 @@ def test_with_init_state(self): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 1]]}, + gates={0: [{0: 1, 1: 1}]}, n_samples=100, key=jax.random.PRNGKey(42), init_state_elems=state_elems, @@ -926,7 +922,7 @@ def test_non_uniform_dims_builds_and_evaluates(self): config = QuditCircuitConfig( dims=[2, 3], n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 2]], 2: [[1, 1]]}, + gates={0: [{0: 1}], 1: [{1: 2}], 2: [{0: 1, 1: 1}]}, n_samples=200, key=jax.random.PRNGKey(42), ) @@ -947,7 +943,7 @@ def test_non_uniform_dims_wires_subset(self): config = QuditCircuitConfig( dims=[2, 3, 4], n_qudits=3, - gates={0: [[1, 0, 0]], 1: [[0, 1, 0]], 2: [[0, 0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}], 2: [{2: 1}]}, n_samples=200, key=jax.random.PRNGKey(0), ) @@ -982,7 +978,7 @@ def _make_config(phase_fn): return QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]], 2: [[1, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]}, n_samples=200, key=jax.random.PRNGKey(42), phase_fn=phase_fn, @@ -1010,7 +1006,7 @@ def test_nontrivial_phase_layer_changes_loss(self): def test_phase_layer_matches_exact_mmd(self): """The estimator with a phase layer must be unbiased w.r.t. exact MMD^2.""" - gates = {0: [[1, 0]], 1: [[0, 1]], 2: [[1, 1]]} + gates = {0: [{0: 1}], 1: [{1: 1}], 2: [{0: 1, 1: 1}]} params = [0.37, 0.95, 0.73] n_qudits, d, bandwidth, graph_type = 2, 3, 0.5, "cycle" @@ -1120,7 +1116,7 @@ def phase_fn(xi, z): config = QuditCircuitConfig( dims=3, n_qudits=2, - gates={0: [[1, 0]], 1: [[0, 1]]}, + gates={0: [{0: 1}], 1: [{1: 1}]}, n_samples=200, key=jax.random.PRNGKey(0), phase_fn=phase_fn, diff --git a/pennylane/labs/tests/tcdq/test_qudit_expval.py b/pennylane/labs/tests/tcdq/test_qudit_expval.py index 2e67f47c67a..b277e05d4d3 100644 --- a/pennylane/labs/tests/tcdq/test_qudit_expval.py +++ b/pennylane/labs/tests/tcdq/test_qudit_expval.py @@ -25,7 +25,8 @@ from pennylane.labs.tcdq.qudit_expval_functions import ( QuditCircuitConfig, _dims_to_numpy, - _parse_qudit_generator_dict, + _gates_to_dense_generators, + _parse_qudit_gate_dict, build_qudit_expval_func, ) @@ -37,10 +38,10 @@ def _build_qudit_expval_func_exact(config): """Build a brute-force reference evaluator by summing over all basis states.""" - generators, param_map = _parse_qudit_generator_dict(config.gates, config.n_qudits) - dims = _dims_to_numpy(config.dims, config.n_qudits) # (n_qudits,) dims_f = jnp.asarray(dims, dtype=jnp.float32) # (n_qudits,) + generators, param_map = _gates_to_dense_generators(config.gates, config.n_qudits, dims) + generators, param_map = jnp.array(generators), jnp.array(param_map) all_states = jnp.array( list(itertools.product(*(range(int(d_j)) for d_j in dims))), @@ -232,13 +233,30 @@ def circuit(): return float(circuit()) +def _dense_rows_to_sparse(generators_array): + """Convert dense generator rows into the sparse ``{qudit: power}`` gate format.""" + return { + i: [{j: int(p) for j, p in enumerate(gen) if int(p) != 0}] + for i, gen in enumerate(generators_array) + } + + +def _sparse_gates_to_rows(gates, n_qudits): + """Expand a list of sparse ``{qudit: power}`` gates into dense generator rows.""" + rows = np.zeros((len(gates), n_qudits), dtype=int) + for i, gate in enumerate(gates): + for qudit, power in gate.items(): + rows[i, qudit] = power + return rows + + def _make_config_one_param_per_gate( d, n, generators_array, thetas, l_vecs, m_vecs, n_samples=NUM_SAMPLES, key=None ): """Build a QuditCircuitConfig with one unique parameter per gate.""" if key is None: key = jax.random.PRNGKey(0) - gates = {i: [list(gen)] for i, gen in enumerate(generators_array)} + gates = _dense_rows_to_sparse(generators_array) return QuditCircuitConfig( dims=d, n_qudits=n, @@ -256,25 +274,25 @@ class TestQuditExpvalVsPennyLane: "n, generators, thetas, l_vecs, m_vecs", [ # Single qubit, X observable - (1, [[1]], [0.37], [[1]], [[0]]), + (1, [{0: 1}], [0.37], [[1]], [[0]]), # Single qubit, Z observable - (1, [[1]], [0.7], [[0]], [[1]]), + (1, [{0: 1}], [0.7], [[0]], [[1]]), # Single qubit, Y observable (D(1,1) = Y) - (1, [[1]], [0.5], [[1]], [[1]]), + (1, [{0: 1}], [0.5], [[1]], [[1]]), # Single qubit, identity observable - (1, [[1]], [0.9], [[0]], [[0]]), + (1, [{0: 1}], [0.9], [[0]], [[0]]), # Two qubits, X0 X1 - (2, [[1, 0], [1, 1]], [0.5, 0.2], [[1, 1]], [[0, 0]]), + (2, [{0: 1}, {0: 1, 1: 1}], [0.5, 0.2], [[1, 1]], [[0, 0]]), # Two qubits, X0 Z1 - (2, [[1, 0], [0, 1]], [0.3, 0.6], [[1, 0]], [[0, 1]]), + (2, [{0: 1}, {1: 1}], [0.3, 0.6], [[1, 0]], [[0, 1]]), # Two qubits, Y0 Y1 (sign = +1 since (-1)^2 = +1) - (2, [[1, 1]], [0.4], [[1, 1]], [[1, 1]]), + (2, [{0: 1, 1: 1}], [0.4], [[1, 1]], [[1, 1]]), # Two qubits, Z0 Z1; two-body gate - (2, [[1, 1]], [0.2], [[0, 0]], [[1, 1]]), + (2, [{0: 1, 1: 1}], [0.2], [[0, 0]], [[1, 1]]), # Three qubits, batch of observables ( 3, - [[1, 0, 0], [0, 1, 0], [0, 0, 1]], + [{0: 1}, {1: 1}, {2: 1}], [0.1, 0.2, 0.3], [[1, 0, 0], [0, 1, 0], [1, 1, 0]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]], @@ -283,7 +301,7 @@ class TestQuditExpvalVsPennyLane: ) def test_matches_pennylane(self, n, generators, thetas, l_vecs, m_vecs): """Qudit expectation values must match exact PennyLane simulation for d=2.""" - generators_arr = np.array(generators) + generators_arr = _sparse_gates_to_rows(generators, n) thetas_arr = np.array(thetas) l_arr = np.array(l_vecs) m_arr = np.array(m_vecs) @@ -303,54 +321,117 @@ def test_matches_pennylane(self, n, generators, thetas, l_vecs, m_vecs): @pytest.mark.parametrize( - "circuit_def, n_qudits, expected_generators, expected_param_map", + "circuit_def, n_qudits, dims, expected_groups", [ - ({0: [[1, 0, 2]]}, 3, [[1, 0, 2]], [0]), - ({0: [[1, 0]], 1: [[0, 2], [1, 1]]}, 2, [[1, 0], [0, 2], [1, 1]], [0, 1, 1]), - ({}, 2, np.zeros((0, 2), dtype=int), []), - ({3: [[0, 1]], 0: [[2, 0]]}, 2, [[2, 0], [0, 1]], [0, 3]), + # one weight-2 gate + ({0: [{0: 1, 2: 2}]}, 3, 3, {2: ([[0, 2]], [[1, 2]], [0])}), + # mixed weights, shared parameter, dict keys given out of order + ( + {0: [{0: 1}], 1: [{1: 2}, {1: 1, 0: 1}]}, + 2, + 3, + {1: ([[0], [1]], [[1], [2]], [0, 1]), 2: ([[0, 1]], [[1, 1]], [1])}, + ), + # empty circuit + ({}, 2, 3, {}), + # parameter keys sorted, not insertion order + ({3: [{1: 1}], 0: [{0: 2}]}, 2, 3, {1: ([[0], [1]], [[2], [1]], [0, 3])}), + # powers reduced modulo the local dimension; zero-power entries dropped + ({0: [{0: 4, 1: 3}]}, 2, 3, {1: ([[0]], [[1]], [0])}), + # identity gate (all powers vanish mod d) is dropped entirely + ({0: [{0: 3}], 1: [{1: 1}]}, 2, 3, {1: ([[1]], [[1]], [1])}), + # negative qudit indices count from the end + ({0: [{-1: 1, 0: 1}]}, 3, 2, {2: ([[0, 2]], [[1, 1]], [0])}), + # per-qudit dimensions + ({0: [{0: 1, 1: 5}]}, 2, [2, 5], {1: ([[0]], [[1]], [0])}), ], ) -def test_parse_qudit_generator_dict(circuit_def, n_qudits, expected_generators, expected_param_map): - """_parse_qudit_generator_dict should produce the correct generator matrix and param map.""" - generators, param_map = _parse_qudit_generator_dict(circuit_def, n_qudits) +def test_parse_qudit_gate_dict(circuit_def, n_qudits, dims, expected_groups): + """_parse_qudit_gate_dict should bucket gates by weight with sorted sparse supports.""" + groups = _parse_qudit_gate_dict(circuit_def, n_qudits, _dims_to_numpy(dims, n_qudits)) + + assert [g.omega for g in groups] == sorted(expected_groups) + for group in groups: + exp_supports, exp_powers, exp_pidx = expected_groups[group.omega] + assert isinstance(group.supports, np.ndarray) + assert isinstance(group.powers, np.ndarray) + assert isinstance(group.param_indices, jnp.ndarray) + assert group.supports.shape == (len(exp_pidx), group.omega) + assert group.powers.shape == (len(exp_pidx), group.omega) + np.testing.assert_array_equal(group.supports, np.array(exp_supports)) + np.testing.assert_array_equal(group.powers, np.array(exp_powers)) + np.testing.assert_array_equal(group.param_indices, np.array(exp_pidx)) - assert isinstance(generators, jnp.ndarray) - assert isinstance(param_map, jnp.ndarray) +@pytest.mark.parametrize( + "circuit_def, n_qudits, dims, expected_generators, expected_param_map", + [ + ({0: [{0: 1, 2: 2}]}, 3, 3, [[1, 0, 2]], [0]), + ({0: [{0: 1}], 1: [{1: 2}, {0: 1, 1: 1}]}, 2, 3, [[1, 0], [0, 2], [1, 1]], [0, 1, 1]), + ({}, 2, 3, np.zeros((0, 2), dtype=int), []), + ({3: [{1: 1}], 0: [{0: 2}]}, 2, 3, [[2, 0], [0, 1]], [0, 3]), + # identity gates are kept so param_map stays aligned with the input + ({0: [{0: 3}], 1: [{1: 1}]}, 2, 3, [[0, 0], [0, 1]], [0, 1]), + ], +) +def test_gates_to_dense_generators( + circuit_def, n_qudits, dims, expected_generators, expected_param_map +): + """The dense expansion used by brute-force references must match the sparse input.""" + generators, param_map = _gates_to_dense_generators( + circuit_def, n_qudits, _dims_to_numpy(dims, n_qudits) + ) expected_generators = np.array(expected_generators) - expected_param_map = np.array(expected_param_map) - assert generators.shape == expected_generators.shape - assert param_map.shape == expected_param_map.shape - assert np.allclose(generators, expected_generators) - assert np.allclose(param_map, expected_param_map) + np.testing.assert_array_equal(generators, expected_generators) + np.testing.assert_array_equal(param_map, np.array(expected_param_map)) + + +def test_parse_qudit_gate_dict_rejects_dense_vectors(): + """A dense length-n_qudits list is no longer a valid gate.""" + with pytest.raises(TypeError, match="dict mapping qudit index"): + _parse_qudit_gate_dict({0: [[1, 0, 2]]}, 3, _dims_to_numpy(3, 3)) -def test_parse_qudit_generator_dict_wrong_length(): - """Generator with wrong length should raise ValueError.""" - with pytest.raises(ValueError, match="length"): - _parse_qudit_generator_dict({0: [[1, 2]]}, n_qudits=3) +@pytest.mark.parametrize("bad_index", [3, 7, -4]) +def test_parse_qudit_gate_dict_index_out_of_range(bad_index): + """Qudit indices outside [-n_qudits, n_qudits) raise IndexError.""" + with pytest.raises(IndexError, match="out of range"): + _parse_qudit_gate_dict({0: [{bad_index: 1}]}, 3, _dims_to_numpy(3, 3)) + + +def test_build_qudit_expval_func_rejects_dense_gates(): + """The public entry point surfaces the TypeError for dense gates.""" + config = QuditCircuitConfig( + dims=3, + n_qudits=2, + gates={0: [[1, 0]]}, + observables=(np.array([[1, 0]]), np.array([[0, 0]])), + n_samples=16, + key=jax.random.PRNGKey(0), + ) + with pytest.raises(TypeError, match="dict mapping qudit index"): + build_qudit_expval_func(config) @pytest.mark.parametrize( "n, thetas, generators, l, m", [ # Single qubit, Z observable - (1, [0.7], [[1]], [0], [1]), + (1, [0.7], [{0: 1}], [0], [1]), # Single qubit, X observable - (1, [0.37], [[1]], [1], [0]), + (1, [0.37], [{0: 1}], [1], [0]), # Two qubits, two gates, X0 Z1 - (2, [0.3, 0.6], [[1, 0], [0, 1]], [1, 0], [0, 1]), + (2, [0.3, 0.6], [{0: 1}, {1: 1}], [1, 0], [0, 1]), # Two qubits, entangling gate, Y0 Y1 - (2, [0.4], [[1, 1]], [1, 1], [1, 1]), + (2, [0.4], [{0: 1, 1: 1}], [1, 1], [1, 1]), # Three qubits, three gates, Z0 I1 X2 - (3, [0.1, 0.2, 0.3], [[1, 0, 0], [0, 1, 0], [0, 0, 1]], [0, 0, 1], [1, 0, 0]), + (3, [0.1, 0.2, 0.3], [{0: 1}, {1: 1}, {2: 1}], [0, 0, 1], [1, 0, 0]), ], ) def test_qudit_expval_exact_matches_pennylane(n, thetas, generators, l, m): """Test _build_qudit_expval_func_exact against PennyLane for d=2.""" - generators = np.array(generators) + generators = _sparse_gates_to_rows(generators, n) thetas = np.array(thetas) l_vecs = np.array([l]) m_vecs = np.array([m]) @@ -372,26 +453,26 @@ class TestQuditExpvalBatchedVsExact: @pytest.mark.parametrize( "d, n, generators, thetas, l_vecs, m_vecs", [ - (2, 1, [[1]], [0.5], [[1]], [[0]]), - (2, 1, [[1]], [0.3], [[0]], [[1]]), - (2, 2, [[1, 0], [1, 1]], [0.5, 0.2], [[1, 1]], [[0, 0]]), + (2, 1, [{0: 1}], [0.5], [[1]], [[0]]), + (2, 1, [{0: 1}], [0.3], [[0]], [[1]]), + (2, 2, [{0: 1}, {0: 1, 1: 1}], [0.5, 0.2], [[1, 1]], [[0, 0]]), ( 2, 2, - [[1, 0], [0, 1]], + [{0: 1}, {1: 1}], [0.4, 0.6], [[1, 0], [0, 1], [1, 1]], [[0, 0], [0, 0], [0, 0]], ), - (3, 1, [[1]], [0.42], [[1]], [[0]]), - (3, 1, [[2]], [0.3], [[1]], [[1]]), - (3, 2, [[1, 0], [0, 2]], [0.5, 0.2], [[1, 1]], [[0, 1]]), - (4, 2, [[1, 2], [3, 1]], [0.3, 0.7], [[0, 0]], [[0, 0]]), + (3, 1, [{0: 1}], [0.42], [[1]], [[0]]), + (3, 1, [{0: 2}], [0.3], [[1]], [[1]]), + (3, 2, [{0: 1}, {1: 2}], [0.5, 0.2], [[1, 1]], [[0, 1]]), + (4, 2, [{0: 1, 1: 2}, {0: 3, 1: 1}], [0.3, 0.7], [[0, 0]], [[0, 0]]), ], ) def test_matches_exact(self, d, n, generators, thetas, l_vecs, m_vecs): """Batched Monte Carlo must agree with the exact qudit expval within sampling noise.""" - generators_arr = np.array(generators) + generators_arr = _sparse_gates_to_rows(generators, n) thetas_arr = np.array(thetas) l_arr = np.array(l_vecs) m_arr = np.array(m_vecs) @@ -424,15 +505,15 @@ class TestQuditExpvalBatchedVsMatrix: @pytest.mark.parametrize( "d, n, generators, thetas, l_vecs, m_vecs", [ - (2, 2, [[1, 0], [0, 1]], [0.3, 0.6], [[1, 0]], [[0, 1]]), - (3, 2, [[1, 0], [0, 2]], [0.5, 0.2], [[1, 1]], [[2, 1]]), - (3, 2, [[1, 2], [0, 2]], [0.1, 0.8], [[1, 1]], [[1, 2]]), - (3, 1, [[2]], [0.3], [[1]], [[1]]), + (2, 2, [{0: 1}, {1: 1}], [0.3, 0.6], [[1, 0]], [[0, 1]]), + (3, 2, [{0: 1}, {1: 2}], [0.5, 0.2], [[1, 1]], [[2, 1]]), + (3, 2, [{0: 1, 1: 2}, {1: 2}], [0.1, 0.8], [[1, 1]], [[1, 2]]), + (3, 1, [{0: 2}], [0.3], [[1]], [[1]]), ], ) def test_matches_matrix_reference(self, d, n, generators, thetas, l_vecs, m_vecs): """Batched Monte Carlo must agree with the dense matrix reference.""" - generators_arr = np.array(generators) + generators_arr = _sparse_gates_to_rows(generators, n) thetas_arr = np.array(thetas) l_arr = np.array(l_vecs) m_arr = np.array(m_vecs) @@ -471,7 +552,7 @@ class TestQuditExpvalNonUniformDims: ( [2, 3], 2, - [[1, 0], [0, 2], [1, 1]], + [{0: 1}, {1: 2}, {0: 1, 1: 1}], [0.5, 0.2, 0.3], [[1, 1], [0, 2]], [[0, 1], [1, 0]], @@ -480,7 +561,7 @@ class TestQuditExpvalNonUniformDims: ( [3, 4], 2, - [[1, 0], [0, 3], [2, 1]], + [{0: 1}, {1: 3}, {0: 2, 1: 1}], [0.4, 0.7, 0.1], [[2, 3], [1, 0]], [[1, 2], [0, 1]], @@ -489,7 +570,7 @@ class TestQuditExpvalNonUniformDims: ( [2, 3, 4], 3, - [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 2, 3]], + [{0: 1}, {1: 1}, {2: 1}, {0: 1, 1: 2, 2: 3}], [0.2, 0.5, 0.3, 0.15], [[1, 0, 0], [0, 2, 3], [1, 1, 1]], [[0, 0, 0], [1, 1, 2], [0, 2, 1]], @@ -498,7 +579,7 @@ class TestQuditExpvalNonUniformDims: ) def test_matches_exact_and_matrix(self, dims, n, generators, thetas, l_vecs, m_vecs): """Batched MC must agree with both the exact and dense-matrix references.""" - generators_arr = np.array(generators) + generators_arr = _sparse_gates_to_rows(generators, n) thetas_arr = np.array(thetas) l_arr = np.array(l_vecs) m_arr = np.array(m_vecs) @@ -667,7 +748,7 @@ def test_empty_gates(self): def test_parameter_broadcasting(self): """Multiple gates sharing a parameter index should all use the same theta.""" d, n = 2, 3 - gates = {0: [[1, 0, 0], [0, 1, 0]], 1: [[0, 0, 1]]} + gates = {0: [{0: 1}, {1: 1}], 1: [{2: 1}]} thetas_unique = np.array([0.5, 0.3]) l_vecs = np.array([[1, 0, 0]]) m_vecs = np.array([[0, 0, 0]]) @@ -839,20 +920,20 @@ def loss(p): "d, n, thetas, generators, l, m", [ # d=2, single qubit, single gate - (2, 1, [0.5], [[1]], [1], [0]), + (2, 1, [0.5], [{0: 1}], [1], [0]), # d=2, two qubits, two-body gate - (2, 2, [0.4], [[1, 1]], [1, 1], [0, 0]), + (2, 2, [0.4], [{0: 1, 1: 1}], [1, 1], [0, 0]), # d=3, single qutrit, single gate - (3, 1, [0.7], [[2]], [1], [2]), + (3, 1, [0.7], [{0: 2}], [1], [2]), # d=3, two qutrits, two gates - (3, 2, [0.3, 0.6], [[1, 0], [0, 2]], [2, 1], [1, 0]), + (3, 2, [0.3, 0.6], [{0: 1}, {1: 2}], [2, 1], [1, 0]), # d=4, single ququart - (4, 1, [1.1], [[3]], [2], [1]), + (4, 1, [1.1], [{0: 3}], [2], [1]), ], ) def test_qudit_expval_batched_matches_exact(d, n, thetas, generators, l, m): """Test build_qudit_expval_func against the exact version for various dimensions.""" - generators = np.array(generators) + generators = _sparse_gates_to_rows(generators, n) thetas = np.array(thetas) l_vecs = np.array([l]) m_vecs = np.array([m]) @@ -890,14 +971,14 @@ class TestQuditExpvalBatchedWithInitState: "d, n, generators, thetas, l_vecs, m_vecs, state_elems, state_amps", [ # d=2, single computational basis state |1> - (2, 1, [[1]], [0.5], [[1]], [[0]], [[1]], [1.0]), + (2, 1, [{0: 1}], [0.5], [[1]], [[0]], [[1]], [1.0]), # d=2, n=2, computational basis state |10> - (2, 2, [[1, 0], [0, 1]], [0.4, 0.6], [[1, 0]], [[0, 0]], [[1, 0]], [1.0]), + (2, 2, [{0: 1}, {1: 1}], [0.4, 0.6], [[1, 0]], [[0, 0]], [[1, 0]], [1.0]), # d=2, n=2, equal superposition (|00> + |11>)/sqrt(2), multiple observables ( 2, 2, - [[1, 0], [0, 1]], + [{0: 1}, {1: 1}], [0.3, 0.5], [[1, 0], [0, 1], [1, 1]], [[0, 0], [0, 0], [0, 0]], @@ -908,7 +989,7 @@ class TestQuditExpvalBatchedWithInitState: ( 2, 2, - [[1, 0], [1, 1]], + [{0: 1}, {0: 1, 1: 1}], [0.4, 0.2], [[1, 1]], [[1, 0]], @@ -916,14 +997,14 @@ class TestQuditExpvalBatchedWithInitState: [1 / np.sqrt(2), 1 / np.sqrt(2)], ), # d=3, single qutrit, basis state |2> - (3, 1, [[1]], [0.42], [[1]], [[0]], [[2]], [1.0]), + (3, 1, [{0: 1}], [0.42], [[1]], [[0]], [[2]], [1.0]), # d=3, n=2, basis state |1, 2> - (3, 2, [[1, 0], [0, 2]], [0.5, 0.2], [[1, 1]], [[0, 1]], [[1, 2]], [1.0]), + (3, 2, [{0: 1}, {1: 2}], [0.5, 0.2], [[1, 1]], [[0, 1]], [[1, 2]], [1.0]), # d=3, n=2, equal superposition (|00> + |12> + |21>)/sqrt(3) ( 3, 2, - [[1, 0], [0, 1]], + [{0: 1}, {1: 1}], [0.3, 0.4], [[1, 0], [0, 1]], [[0, 0], [0, 0]], @@ -934,7 +1015,7 @@ class TestQuditExpvalBatchedWithInitState: ( 3, 2, - [[1, 0]], + [{0: 1}], [0.5], [[1, 0]], [[0, 0]], @@ -945,7 +1026,7 @@ class TestQuditExpvalBatchedWithInitState: ( 2, 2, - [[1, 0], [0, 1]], + [{0: 1}, {1: 1}], [0.3, 0.5], [[1, 1]], [[1, 1]], @@ -956,7 +1037,7 @@ class TestQuditExpvalBatchedWithInitState: ( 3, 2, - [[1, 0], [0, 1]], + [{0: 1}, {1: 1}], [0.3, 0.4], [[1, 2], [0, 1]], [[2, 1], [0, 0]], @@ -964,16 +1045,16 @@ class TestQuditExpvalBatchedWithInitState: [1 / np.sqrt(2), 1 / np.sqrt(2)], ), # d=4, n=2, basis state - (4, 2, [[1, 2], [3, 1]], [0.3, 0.7], [[1, 0]], [[0, 0]], [[0, 1]], [1.0]), + (4, 2, [{0: 1, 1: 2}, {0: 3, 1: 1}], [0.3, 0.7], [[1, 0]], [[0, 0]], [[0, 1]], [1.0]), # d=4, n=1, complex amplitudes with m != 0 - (4, 1, [[2]], [0.5], [[1]], [[2]], [[0], [3]], [1 / np.sqrt(2), 1j / np.sqrt(2)]), + (4, 1, [{0: 2}], [0.5], [[1]], [[2]], [[0], [3]], [1 / np.sqrt(2), 1j / np.sqrt(2)]), ], ) def test_matches_matrix_reference( self, d, n, generators, thetas, l_vecs, m_vecs, state_elems, state_amps ): """Batched Monte Carlo with init state must agree with dense matrix reference.""" - generators_arr = np.array(generators) + generators_arr = _sparse_gates_to_rows(generators, n) thetas_arr = np.array(thetas) l_arr = np.array(l_vecs) m_arr = np.array(m_vecs) @@ -1059,7 +1140,7 @@ def test_config_init_state(self): elems = np.array([[0, 0], [1, 1]]) amps = np.array([1 / np.sqrt(2), 1 / np.sqrt(2)], dtype=complex) - gates = {i: [list(gen)] for i, gen in enumerate(generators)} + gates = _dense_rows_to_sparse(generators) config = QuditCircuitConfig( dims=d, n_qudits=n, @@ -1102,7 +1183,7 @@ def test_runtime_override_takes_precedence(self): runtime_elems = np.array([[1, 1]]) runtime_amps = np.array([1.0 + 0j]) - gates = {i: [list(gen)] for i, gen in enumerate(generators)} + gates = _dense_rows_to_sparse(generators) config = QuditCircuitConfig( dims=d, n_qudits=n, @@ -1324,28 +1405,28 @@ def loss(p): "d, n, thetas, generators, l, m, state_elems, state_amps", [ # d=2, single qubit, computational basis |1> - (2, 1, [0.5], [[1]], [1], [0], [[1]], [1.0]), + (2, 1, [0.5], [{0: 1}], [1], [0], [[1]], [1.0]), # d=2, two qubits, |10> - (2, 2, [0.4, 0.6], [[1, 0], [0, 1]], [1, 0], [0, 0], [[1, 0]], [1.0]), + (2, 2, [0.4, 0.6], [{0: 1}, {1: 1}], [1, 0], [0, 0], [[1, 0]], [1.0]), # d=2, two qubits, superposition (|00> + |11>)/sqrt(2) ( 2, 2, [0.3, 0.5], - [[1, 0], [0, 1]], + [{0: 1}, {1: 1}], [1, 0], [0, 0], [[0, 0], [1, 1]], [1 / np.sqrt(2), 1 / np.sqrt(2)], ), # d=3, single qutrit, |2> with complex amplitude - (3, 1, [0.7], [[2]], [1], [2], [[2]], [1.0 + 0j]), + (3, 1, [0.7], [{0: 2}], [1], [2], [[2]], [1.0 + 0j]), # d=3, two qutrits, superposition (|01> + i|20>)/sqrt(2) ( 3, 2, [0.3], - [[1, 2]], + [{0: 1, 1: 2}], [2, 1], [1, 0], [[0, 1], [2, 0]], @@ -1357,7 +1438,7 @@ def test_qudit_expval_batched_init_state_matches_brute_force( d, n, thetas, generators, l, m, state_elems, state_amps ): """Test build_qudit_expval_func with init_state against the dense matrix reference.""" - generators = np.array(generators) + generators = _sparse_gates_to_rows(generators, n) thetas = np.array(thetas) state_elems = np.array(state_elems) state_amps = np.array(state_amps, dtype=complex) @@ -1419,7 +1500,7 @@ def test_phase_layer_default_state(self): phase_params = jnp.array([0.1, 0.5, 2.0, 1.0]) n_samples = 80000 - gates = {i: [list(gen)] for i, gen in enumerate(generators)} + gates = _dense_rows_to_sparse(generators) config = QuditCircuitConfig( dims=d, n_qudits=n, @@ -1465,7 +1546,7 @@ def test_phase_layer_with_init_state(self): state_amps = np.array([1 / np.sqrt(3), 1 / np.sqrt(3), 1 / np.sqrt(3)], dtype=complex) n_samples = 80000 - gates = {i: [list(gen)] for i, gen in enumerate(generators)} + gates = _dense_rows_to_sparse(generators) config = QuditCircuitConfig( dims=d, n_qudits=n, @@ -1515,7 +1596,7 @@ def test_phase_layer_grad(self): m_vecs = np.array([[0, 0]]) phase_params = jnp.array([0.1, 0.5, 2.0]) - gates = {i: [list(gen)] for i, gen in enumerate(generators)} + gates = _dense_rows_to_sparse(generators) config = QuditCircuitConfig( dims=d, n_qudits=n, From 4571323229808436e77d13c4755ca66f666f6108 Mon Sep 17 00:00:00 2001 From: Will Date: Thu, 8 Oct 2026 13:44:13 -0400 Subject: [PATCH 3/7] update utils --- pennylane/labs/tcdq/__init__.py | 2 + pennylane/labs/tcdq/utils.py | 67 +++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/pennylane/labs/tcdq/__init__.py b/pennylane/labs/tcdq/__init__.py index 96fa31272fe..ab425fd4a76 100644 --- a/pennylane/labs/tcdq/__init__.py +++ b/pennylane/labs/tcdq/__init__.py @@ -330,6 +330,7 @@ def loss_fn(current_params): create_lattice_gates, create_local_gates, create_random_gates, + create_random_qudit_gates, generate_pauli_observables, ) @@ -351,5 +352,6 @@ def loss_fn(current_params): "create_lattice_gates", "create_local_gates", "create_random_gates", + "create_random_qudit_gates", "generate_pauli_observables", ] diff --git a/pennylane/labs/tcdq/utils.py b/pennylane/labs/tcdq/utils.py index dd741ae623a..405133b9e1a 100644 --- a/pennylane/labs/tcdq/utils.py +++ b/pennylane/labs/tcdq/utils.py @@ -17,6 +17,7 @@ :class:`~pennylane.labs.tcdq.CircuitConfig`. """ +from collections.abc import Sequence from itertools import combinations import networkx as nx @@ -200,3 +201,69 @@ def generate_pauli_observables( observables.append(obs_row) return observables + + +def create_random_qudit_gates( + n_gates: int, + dims: Sequence[int], + min_weight: int = 1, + max_weight: int = 2, + seed: int = None, +) -> dict[int, list[dict[int, int]]]: + r"""Generate a sparse qudit gate dictionary with randomly chosen supports and powers. + + This is the qudit analogue of :func:`~pennylane.labs.tcdq.create_random_gates`. + Each gate acts on a uniformly random subset of qudits whose size (the gate + *weight*) is drawn uniformly from ``[min_weight, max_weight]``. On every + active qudit :math:`j` the power of :math:`Z` is drawn uniformly from + :math:`\{1, \ldots, d_j - 1\}`, so each gate has exactly the drawn weight. + + Args: + n_gates (int): Number of gates (and trainable parameters) to generate. + dims (Sequence[int]): Local dimension :math:`d_j \geq 2` of each qudit. The + number of qudits is ``len(dims)``. + min_weight (int): Minimum number of active qudits per gate. Defaults to ``1``. + max_weight (int): Maximum number of active qudits per gate. Defaults to ``2``. + seed (int): Random seed for reproducibility. Defaults to ``None``. + + Returns: + dict[int, list[dict[int, int]]]: A gate dictionary suitable for + :class:`~pennylane.labs.tcdq.QuditCircuitConfig`. Parameter index ``i`` + maps to a single sparse gate ``{qudit_index: z_power}`` with sorted keys. + + Raises: + ValueError: If ``n_gates`` is negative, if any dimension is smaller than 2, or if + the weight bounds do not satisfy ``1 <= min_weight <= max_weight <= len(dims)``. + + **Example** + + >>> from pennylane.labs.tcdq import create_random_qudit_gates + >>> gates = create_random_qudit_gates(n_gates=4, dims=[3, 3, 4, 2], max_weight=3, seed=0) + >>> len(gates) + 4 + >>> all(1 <= len(gate) <= 3 for (gate,) in gates.values()) + True + """ + dims = [int(d) for d in dims] + n_qudits = len(dims) + + if n_gates < 0: + raise ValueError(f"n_gates must be non-negative, got {n_gates}.") + if any(d < 2 for d in dims): + raise ValueError(f"All qudit dimensions must be at least 2, got dims={dims}.") + if not 1 <= min_weight <= max_weight <= n_qudits: + raise ValueError( + "Weights must satisfy 1 <= min_weight <= max_weight <= len(dims); " + f"got min_weight={min_weight}, max_weight={max_weight}, len(dims)={n_qudits}." + ) + + rng = np.random.default_rng(seed) + gates_dict = {} + + from tqdm import tqdm + for i in tqdm(range(n_gates)): + weight = int(rng.integers(min_weight, max_weight + 1)) + support = sorted(int(q) for q in rng.choice(n_qudits, size=weight, replace=False)) + gates_dict[i] = [{q: int(rng.integers(1, dims[q])) for q in support}] + + return gates_dict From 02ae452249b63a8b66e2624781f680984046ca57 Mon Sep 17 00:00:00 2001 From: Will Date: Thu, 8 Oct 2026 15:48:55 -0400 Subject: [PATCH 4/7] memory control --- pennylane/labs/tcdq/mmd_loss_hw.py | 19 +- pennylane/labs/tcdq/qudit_expval_functions.py | 316 +++++++++--------- .../labs/tests/tcdq/test_qudit_expval.py | 106 ++++++ 3 files changed, 282 insertions(+), 159 deletions(-) diff --git a/pennylane/labs/tcdq/mmd_loss_hw.py b/pennylane/labs/tcdq/mmd_loss_hw.py index 1d48bf98205..083e16879d8 100644 --- a/pennylane/labs/tcdq/mmd_loss_hw.py +++ b/pennylane/labs/tcdq/mmd_loss_hw.py @@ -123,7 +123,7 @@ def _marginal_probs(d: int, bandwidth: float, graph_type: str) -> jnp.ndarray: raise ValueError(f"Unknown graph_type {graph_type!r}; use 'cycle' or 'complete'.") -def _sample_fourier_indices( # pylint: disable=too-many-arguments +def _sample_fourier_indices( # pylint: disable=too-many-arguments,too-many-locals key: ArrayLike, n_ops: int, n_qudits: int, @@ -146,13 +146,20 @@ def _sample_fourier_indices( # pylint: disable=too-many-arguments Integer array of shape ``(n_ops, n_qudits)``; column ``i`` has entries in :math:`\\{0, \\ldots, d_i-1\\}`. """ + # Group the visible wires by local dimension so that every group is sampled with a + # single ``jax.random.choice`` and written with a single scatter. Looping over wires + # instead would trace one choice and one full-width scatter per qudit, making the + # jitted graph (and its compile time and memory) grow with ``n_qudits``. + wires_by_dim: dict[int, list[int]] = {} + for wire in wire_tuple: + wires_by_dim.setdefault(int(dims[wire]), []).append(int(wire)) + all_obs = jnp.zeros((n_ops, n_qudits), dtype=jnp.int32) - keys = jax.random.split(key, len(wire_tuple)) if wire_tuple else [] - for col_key, wire in zip(keys, wire_tuple): - d_i = int(dims[wire]) + keys = jax.random.split(key, len(wires_by_dim)) if wires_by_dim else [] + for group_key, (d_i, wires) in zip(keys, sorted(wires_by_dim.items())): marginal = _marginal_probs(d_i, bandwidth, graph_type) - col = jax.random.choice(col_key, d_i, shape=(n_ops,), p=marginal) - all_obs = all_obs.at[:, wire].set(col.astype(jnp.int32)) + cols = jax.random.choice(group_key, d_i, shape=(n_ops, len(wires)), p=marginal) + all_obs = all_obs.at[:, jnp.asarray(wires, dtype=jnp.int32)].set(cols.astype(jnp.int32)) return all_obs diff --git a/pennylane/labs/tcdq/qudit_expval_functions.py b/pennylane/labs/tcdq/qudit_expval_functions.py index 17bda6ad66e..99bd9ef39c5 100644 --- a/pennylane/labs/tcdq/qudit_expval_functions.py +++ b/pennylane/labs/tcdq/qudit_expval_functions.py @@ -104,6 +104,11 @@ class QuditCircuitConfig: # pylint: disable=too-many-instance-attributes :math:`D'(\mathbf{\xi})` is given by a ``Callable`` with signature ``(params: ArrayLike, z: ArrayLike) -> scalar`` where ``z`` is a dit-string of shape ``(n_qudits, )`` with entries in :math:`\{0, \dots, d-1\}` and ``params`` has shape matching :math:`\mathbf{\xi}`. + max_memory (float): Upper bound, in gigabytes (1 GB = 1024**3 bytes), on the temporary + memory used by the phase-difference computation. Within each weight group the gates + are processed in blocks whose size is derived from this budget together with the + gate weight, the number of samples and the number of observables, so peak memory no + longer grows with the number of gates. Defaults to ``1.0``. **Example** @@ -141,6 +146,8 @@ class QuditCircuitConfig: # pylint: disable=too-many-instance-attributes init_state_amps: ArrayLike | None = None #: Learnable phase layer phase_fn: Callable | None = None + #: Memory budget in GB for the blocked phase-difference computation. + max_memory: float = 1.0 def _dims_to_numpy(dims: int | Sequence[int], n_qudits: int) -> np.ndarray: @@ -322,136 +329,43 @@ def _compute_qudit_samples( return jax.random.randint(key, shape=(num_samples, n_qudits), minval=0, maxval=maxval) -class WeightGroupData(NamedTuple): - """Precomputed factor matrices for gates sharing the same weight (number of active qudits). +_BYTES_PER_GB = 1024**3 +_FLOAT_BYTES = 4 - Gates are grouped by weight :math:`\\omega` (number of non-zero entries in - the generator vector) so that the :math:`2^\\omega`-term angle-addition - expansion can be vectorised over gates within each group. - - Args: - param_indices (jnp.ndarray): Maps each gate in this group to its parameter - index in the global ``gates_params`` array, shape ``(n_gates,)``. - samples_matrices (list[jnp.ndarray]): :math:`2^\\omega` matrices of shape - ``(n_gates, n_samples)`` giving the sample-side factor for each - angle-addition term. - obs_matrices (list[jnp.ndarray]): :math:`2^\\omega` matrices of shape - ``(n_gates, n_obs)`` giving the observable-side factor for each - angle-addition term. - """ - - #: Maps each gate to its parameter index, shape ``(n_gates,)``. - param_indices: jnp.ndarray - #: Sample-side factor matrices for each angle-addition term. - samples_matrices: list[jnp.ndarray] - #: Observable-side factor matrices for each angle-addition term. - obs_matrices: list[jnp.ndarray] +def _memory_budget_bytes(max_memory: float) -> int: + """Convert a ``max_memory`` budget in GB to an integer number of bytes.""" + return int(max_memory * _BYTES_PER_GB) -def _gather_support_values( - vectors: ArrayLike, supports: np.ndarray, target_dim: int, n_gates: int, omega: int -) -> jnp.ndarray: - """Extract values at the active qudit positions for every gate, for every vector. - - Each gate acts on ``omega`` qudits (its *support*). Given a batch of - full-length vectors (e.g. Monte Carlo samples or observable ``l``-vectors), - this function selects only the entries at each gate's support positions and - arranges them into shape ``(n_gates, omega, target_dim)`` so downstream - trigonometric computations can be vectorised over gates and positions. - Args: - vectors (ArrayLike): Input array of shape ``(target_dim, n_qudits)`` — - either the Monte Carlo samples (``target_dim = n_samples``) or the - observable ``l``-vectors (``target_dim = n_obs``). - supports (np.ndarray): Active qudit indices for each gate, shape - ``(n_gates, omega)``. - target_dim (int): Number of vectors (rows in ``vectors``). - n_gates (int): Number of gates in this weight group. - omega (int): Number of active qudits per gate. +def _qudit_phase_bytes_per_gate(omega: int, n_samples: int, n_obs: int) -> int: + """Approximate temporary bytes needed per gate of weight ``omega`` inside one block. - Returns: - jnp.ndarray: Values at support positions, shape ``(n_gates, omega, target_dim)``. + For every active position the block holds the angle, cosine and sine arrays on both + the sample side (``n_samples`` columns) and the observable side (``n_obs`` columns), + plus the :math:`2^\\omega` angle-addition products on each side. """ - flat_supports = supports.reshape(-1) - return ( - jnp.array(vectors)[:, flat_supports].reshape(target_dim, n_gates, omega).transpose(1, 2, 0) - ) + per_column = _FLOAT_BYTES * (3 * omega + 2**omega) + return per_column * (n_samples + n_obs) -def _compute_trigonometric_building_blocks( - gate_vals: np.ndarray, - z_at_support: jnp.ndarray, - l_at_support: jnp.ndarray, - d_at_support: jnp.ndarray, -) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Compute (state_cos, state_sin, obs_cos, obs_sin) trig factors over the gate support.""" - g = jnp.array(gate_vals, dtype=jnp.float32)[:, :, jnp.newaxis] - d_s = jnp.asarray(d_at_support, dtype=jnp.float32)[:, :, jnp.newaxis] - angle_z = 2 * jnp.pi * g * z_at_support.astype(jnp.float32) / d_s + jnp.pi / 4 - angle_l = 2 * jnp.pi * g * l_at_support.astype(jnp.float32) / d_s - return ( - jnp.sqrt(2.0) * jnp.cos(angle_z), - jnp.sqrt(2.0) * jnp.sin(angle_z), - jnp.cos(angle_l), - jnp.sin(angle_l), - ) - - -def _expand_angle_addition( - state_cos: jnp.ndarray, - state_sin: jnp.ndarray, - obs_cos: jnp.ndarray, - obs_sin: jnp.ndarray, -) -> tuple[list[jnp.ndarray], list[jnp.ndarray]]: - """Enumerate all :math:`2^\\omega` angle-addition terms to build the factor matrices. +def _qudit_phase_block_size(max_memory_bytes: int, omega: int, n_samples: int, n_obs: int) -> int: + """Number of weight-``omega`` gates per block that fits within ``max_memory_bytes``.""" + return max(1, max_memory_bytes // _qudit_phase_bytes_per_gate(omega, n_samples, n_obs)) - Each term corresponds to a binary choice (cos or sin) at each active - qudit position, producing paired sample-side and observable-side factors. - """ - n_gates, omega, num_samples = state_cos.shape - n_obs = obs_cos.shape[2] - state_factors = [state_cos, state_sin] - obs_factors = [obs_cos, obs_sin] - samples_list: list[jnp.ndarray] = [] - obs_list: list[jnp.ndarray] = [] - for sigma in itertools.product([0, 1], repeat=omega): - B = jnp.ones((n_gates, num_samples), dtype=jnp.float32) - C = jnp.ones((n_gates, n_obs), dtype=jnp.float32) - for k, choice in enumerate(sigma): - B *= state_factors[choice][:, k, :] - C *= obs_factors[choice][:, k, :] - samples_list.append(B) - obs_list.append(C) - return samples_list, obs_list +def _obs_phase_matrix( + samples: jnp.ndarray, m_f: jnp.ndarray, l_f: jnp.ndarray, dims: ArrayLike +) -> jnp.ndarray: + """Compute the observable phase matrix. -def _build_weight_group( - group: SparseGateGroup, - samples: jnp.ndarray, - l_vecs: jnp.ndarray, - dims: np.ndarray, -) -> WeightGroupData: - """Precompute the factor matrices for a group of gates with the same weight. - - The supports and powers come straight from the sparse gate description, so - no dense generator matrix is scanned here. + :math:`J[i, j] = \\exp(i\\pi \\sum_k m_{ik} (2 z_{jk} - l_{ik}) / d_k)`. """ - n_gates, omega = group.supports.shape - num_samples = samples.shape[0] - n_obs = l_vecs.shape[0] - d_at_support = np.asarray(dims)[group.supports] # (n_gates, omega) - - z_at_support = _gather_support_values(samples, group.supports, num_samples, n_gates, omega) - l_at_support = _gather_support_values(l_vecs, group.supports, n_obs, n_gates, omega) - - state_cos, state_sin, obs_cos, obs_sin = _compute_trigonometric_building_blocks( - group.powers, z_at_support, l_at_support, d_at_support - ) - samples_matrices, obs_matrices = _expand_angle_addition(state_cos, state_sin, obs_cos, obs_sin) - return WeightGroupData( - param_indices=group.param_indices, - samples_matrices=samples_matrices, - obs_matrices=obs_matrices, + s_f = samples.astype(jnp.float32) + inv_d = (1.0 / jnp.asarray(dims, dtype=jnp.float32))[jnp.newaxis, :] # (1, n_qudits) + m_scaled = m_f * inv_d # (n_obs, n_qudits) + return jnp.exp( + 1j * jnp.pi * (2 * m_scaled @ s_f.T - jnp.sum(m_scaled * l_f, axis=1, keepdims=True)) ) @@ -462,53 +376,145 @@ class _PrecomputedObsData(NamedTuple): n_obs: int l_f: jnp.ndarray m_f: jnp.ndarray - weight_data: list obs_phase_matrix: jnp.ndarray -def _obs_phase_matrix( - samples: jnp.ndarray, m_f: jnp.ndarray, l_f: jnp.ndarray, dims: ArrayLike +# pylint: disable=too-many-arguments +def _group_block_contribution( + theta_ext: jnp.ndarray, + samples_t: jnp.ndarray, + l_t: jnp.ndarray, + dims_f: jnp.ndarray, + supports: jnp.ndarray, + powers: jnp.ndarray, + param_indices: jnp.ndarray, ) -> jnp.ndarray: - """Compute the observable phase matrix. + """Phase-difference contribution of one block of equal-weight gates. - :math:`J[i, j] = \\exp(i\\pi \\sum_k m_{ik} (2 z_{jk} - l_{ik}) / d_k)`. + For every gate :math:`g` in the block the eigenvalue factor + :math:`\\phi_g(z) = \\prod_k \\sqrt{2}\\cos(2\\pi g_k z_k / d_k + \\pi/4)` is expanded with the + angle-addition formula into :math:`2^\\omega` products of sample-side and + observable-side factors, so that :math:`\\sum_g \\theta_g [\\phi_g(z) - \\phi_g(z - l)]` is + assembled from matrix products of shape ``(n_obs, block) @ (block, n_samples)``. + + Args: + theta_ext (jnp.ndarray): Gate parameters with a trailing zero appended, shape + ``(n_params + 1,)``. Padded gates index the trailing zero and contribute nothing. + samples_t (jnp.ndarray): Transposed samples, shape ``(n_qudits, n_samples)``, float. + l_t (jnp.ndarray): Transposed observable ``l`` vectors, shape ``(n_qudits, n_obs)``, float. + dims_f (jnp.ndarray): Per-qudit dimensions, shape ``(n_qudits,)``, float. + supports (jnp.ndarray): Active qudit indices of the block, shape ``(block, omega)``. + powers (jnp.ndarray): :math:`Z` powers of the block, shape ``(block, omega)``. + param_indices (jnp.ndarray): Parameter index of each gate, shape ``(block,)``. + + Returns: + jnp.ndarray: Accumulated phase differences, shape ``(n_obs, n_samples)``. """ - s_f = samples.astype(jnp.float32) - inv_d = (1.0 / jnp.asarray(dims, dtype=jnp.float32))[jnp.newaxis, :] # (1, n_qudits) - m_scaled = m_f * inv_d # (n_obs, n_qudits) - return jnp.exp( - 1j * jnp.pi * (2 * m_scaled @ s_f.T - jnp.sum(m_scaled * l_f, axis=1, keepdims=True)) + omega = supports.shape[1] + theta = theta_ext[param_indices] # (block,) + + g = powers.astype(jnp.float32)[:, :, jnp.newaxis] # (block, omega, 1) + d_s = dims_f[supports][:, :, jnp.newaxis] # (block, omega, 1) + z_at_support = jnp.take(samples_t, supports, axis=0) # (block, omega, n_samples) + l_at_support = jnp.take(l_t, supports, axis=0) # (block, omega, n_obs) + + angle_z = 2 * jnp.pi * g * z_at_support / d_s + jnp.pi / 4 + angle_l = 2 * jnp.pi * g * l_at_support / d_s + state_factors = (jnp.sqrt(2.0) * jnp.cos(angle_z), jnp.sqrt(2.0) * jnp.sin(angle_z)) + obs_factors = (jnp.cos(angle_l), jnp.sin(angle_l)) + + accumulated = None + for sigma in itertools.product([0, 1], repeat=omega): + B_sigma = state_factors[sigma[0]][:, 0, :] + C_sigma = obs_factors[sigma[0]][:, 0, :] + for k in range(1, omega): + B_sigma = B_sigma * state_factors[sigma[k]][:, k, :] + C_sigma = C_sigma * obs_factors[sigma[k]][:, k, :] + if accumulated is None: # sigma == (0, ..., 0): B_sigma is the unshifted phi_g(z) + accumulated = (theta @ B_sigma)[jnp.newaxis, :] + accumulated = accumulated - (C_sigma.T * theta) @ B_sigma + + return accumulated + + +def _group_phase_diffs( + theta_ext: jnp.ndarray, + samples_t: jnp.ndarray, + l_t: jnp.ndarray, + dims_f: jnp.ndarray, + group: SparseGateGroup, + block_size: int, +) -> jnp.ndarray: + """Phase differences of one weight group, processed in memory-bounded blocks. + + Each block is wrapped in :func:`jax.checkpoint` so that reverse-mode + differentiation recomputes the block's trigonometric factors instead of + storing them, keeping gradient memory bounded as well. + """ + n_gates, omega = group.supports.shape + supports = jnp.asarray(group.supports, dtype=jnp.int32) + powers = jnp.asarray(group.powers, dtype=jnp.int32) + param_indices = jnp.asarray(group.param_indices, dtype=jnp.int32) + + block_fn = jax.checkpoint( + lambda s, p, i: _group_block_contribution(theta_ext, samples_t, l_t, dims_f, s, p, i) ) + if n_gates <= block_size: + return block_fn(supports, powers, param_indices) + + n_blocks = -(-n_gates // block_size) + n_pad = n_blocks * block_size - n_gates + if n_pad: + pad_index = theta_ext.shape[0] - 1 # points at the appended zero parameter + supports = jnp.concatenate([supports, jnp.zeros((n_pad, omega), supports.dtype)]) + powers = jnp.concatenate([powers, jnp.ones((n_pad, omega), powers.dtype)]) + param_indices = jnp.concatenate( + [param_indices, jnp.full((n_pad,), pad_index, param_indices.dtype)] + ) -def _build_all_weight_groups( - gate_groups: list[SparseGateGroup], - samples: jnp.ndarray, - l_vecs: jnp.ndarray, - dims: np.ndarray, -) -> list[WeightGroupData]: - """Build :class:`WeightGroupData` for each sparse gate group.""" - return [_build_weight_group(group, samples, l_vecs, dims) for group in gate_groups] + def accumulate(total, block): + return total + block_fn(*block), None + + carry_dtype = jnp.result_type(theta_ext.dtype, samples_t.dtype) + zero = jnp.zeros((l_t.shape[1], samples_t.shape[1]), dtype=carry_dtype) + total, _ = jax.lax.scan( + accumulate, + zero, + ( + supports.reshape(n_blocks, block_size, omega), + powers.reshape(n_blocks, block_size, omega), + param_indices.reshape(n_blocks, block_size), + ), + ) + return total -# pylint: disable=too-many-arguments def _accumulate_phase_diffs( gates_params: ArrayLike, - weight_data: list[WeightGroupData], - n_obs: int, - n_samples: int, + gate_groups: list[SparseGateGroup], + samples: jnp.ndarray, + l_vecs: jnp.ndarray, + dims: ArrayLike, + max_memory_bytes: int, vmapped_phase_func: Callable | None, phase_fn_params: ArrayLike | None, - samples: ArrayLike, - l_vecs: ArrayLike, ) -> jnp.ndarray: """Assemble the accumulated phase-difference matrix from all weight groups.""" + n_samples = samples.shape[0] + n_obs = l_vecs.shape[0] + theta = jnp.asarray(gates_params) + theta_ext = jnp.concatenate([theta, jnp.zeros((1,), dtype=theta.dtype)]) + samples_t = samples.astype(jnp.float32).T # (n_qudits, n_samples) + l_t = jnp.asarray(l_vecs).astype(jnp.float32).T # (n_qudits, n_obs) + dims_f = jnp.asarray(dims, dtype=jnp.float32) + accumulated = jnp.zeros((n_obs, n_samples)) - for group in weight_data: - theta_w = jnp.asarray(gates_params)[group.param_indices] - accumulated = accumulated + (theta_w @ group.samples_matrices[0])[jnp.newaxis, :] - for B_sigma, C_sigma in zip(group.samples_matrices, group.obs_matrices): - accumulated = accumulated - (C_sigma.T * theta_w) @ B_sigma + for group in gate_groups: + block_size = _qudit_phase_block_size(max_memory_bytes, group.omega, n_samples, n_obs) + accumulated = accumulated + _group_phase_diffs( + theta_ext, samples_t, l_t, dims_f, group, block_size + ) if vmapped_phase_func is not None: accumulated += vmapped_phase_func(phase_fn_params, samples, l_vecs) @@ -645,6 +651,7 @@ def build_qudit_expval_func( # pylint: disable=too-many-statements n = config.n_qudits dims = _dims_to_numpy(config.dims, n) gate_groups = _parse_qudit_gate_dict(config.gates, n, dims) + max_memory_bytes = _memory_budget_bytes(config.max_memory) default_samples = _compute_qudit_samples(config.key, config.n_samples, n, dims) vmapped_phase_func = None @@ -672,7 +679,6 @@ def compute_phase_diff(p_params, sample, l_vec): n_obs=n_obs, l_f=l_f, m_f=m_f, - weight_data=_build_all_weight_groups(gate_groups, default_samples, l_vecs, dims), obs_phase_matrix=_obs_phase_matrix(default_samples, m_f, l_f, dims), ) else: @@ -715,11 +721,10 @@ def qudit_expval_batched( """ if observables is not None: l_vecs = jnp.array(observables[0], dtype=jnp.int32) - n_obs = l_vecs.shape[0] l_f = l_vecs.astype(jnp.float32) m_f = jnp.array(observables[1], dtype=jnp.int32).astype(jnp.float32) elif defaults is not None: - l_vecs, n_obs, l_f, m_f = defaults.l_vecs, defaults.n_obs, defaults.l_f, defaults.m_f + l_vecs, l_f, m_f = defaults.l_vecs, defaults.l_f, defaults.m_f else: raise ValueError( "No observables specified. Provide them in QuditCircuitConfig " @@ -739,13 +744,18 @@ def qudit_expval_batched( ) if use_cached: obs_pm = defaults.obs_phase_matrix - w_data = defaults.weight_data else: obs_pm = _obs_phase_matrix(samples, m_f, l_f, dims) - w_data = _build_all_weight_groups(gate_groups, samples, l_vecs, dims) accumulated_phase_diffs = _accumulate_phase_diffs( - gates_params, w_data, n_obs, _n, vmapped_phase_func, phase_fn_params, samples, l_vecs + gates_params, + gate_groups, + samples, + l_vecs, + dims, + max_memory_bytes, + vmapped_phase_func, + phase_fn_params, ) state_elems = config.init_state_elems if init_state_elems is None else init_state_elems diff --git a/pennylane/labs/tests/tcdq/test_qudit_expval.py b/pennylane/labs/tests/tcdq/test_qudit_expval.py index b277e05d4d3..91132e7e3c0 100644 --- a/pennylane/labs/tests/tcdq/test_qudit_expval.py +++ b/pennylane/labs/tests/tcdq/test_qudit_expval.py @@ -23,10 +23,13 @@ import pennylane as qp from pennylane.labs.tcdq.qudit_expval_functions import ( + _BYTES_PER_GB, QuditCircuitConfig, _dims_to_numpy, _gates_to_dense_generators, _parse_qudit_gate_dict, + _qudit_phase_block_size, + _qudit_phase_bytes_per_gate, build_qudit_expval_func, ) @@ -1616,3 +1619,106 @@ def loss(p_params): grad_val = jax.grad(loss)(phase_params) assert grad_val.shape == phase_params.shape assert not jnp.allclose(grad_val, 0.0) + + +class TestQuditPhaseBlocking: + """The memory-bounded blocked phase computation must match the single-block result.""" + + @staticmethod + def _random_gates(n_qudits, n_gates, dims, seed): + rng = np.random.default_rng(seed) + gates = {} + for i in range(n_gates): + weight = int(rng.integers(1, 4)) + support = sorted(int(q) for q in rng.choice(n_qudits, size=weight, replace=False)) + gates[i] = [{q: int(rng.integers(1, dims[q])) for q in support}] + gates[0].append({0: 1, 1: 1}) # two gates sharing one parameter + return gates + + def test_block_size_helper(self): + """Block size is at least one and shrinks with weight, samples and observables.""" + budget = 10 * _qudit_phase_bytes_per_gate(2, 100, 50) + assert _qudit_phase_block_size(budget, 2, 100, 50) == 10 + assert _qudit_phase_block_size(budget, 1, 100, 50) > 10 + assert _qudit_phase_block_size(budget, 3, 100, 50) < 10 + assert _qudit_phase_block_size(budget, 2, 1000, 50) < 10 + assert _qudit_phase_block_size(1, 4, 10_000, 10_000) == 1 + + @pytest.mark.parametrize("dims", [[3] * 6, [2, 3, 4, 5, 3, 2]]) + def test_blocked_matches_single_block(self, dims): + """Forcing many small blocks (with padding) reproduces the single-block values and gradients.""" + n_qudits, n_gates, n_samples, n_obs = len(dims), 23, 256, 7 + gates = self._random_gates(n_qudits, n_gates, dims, seed=0) + rng = np.random.default_rng(1) + dims_np = _dims_to_numpy(dims, n_qudits) + l_vecs = np.stack([rng.integers(0, dims_np) for _ in range(n_obs)]) + m_vecs = np.stack([rng.integers(0, dims_np) for _ in range(n_obs)]) + elems = np.stack([rng.integers(0, dims_np) for _ in range(3)]) + amps = rng.normal(size=3) + 1j * rng.normal(size=3) + amps = amps / np.linalg.norm(amps) + + def phase_fn(p, z): + return jnp.sum(p[:n_qudits] * z.astype(jnp.float32)) + p[-1] + + common = { + "dims": dims, + "n_qudits": n_qudits, + "gates": gates, + "observables": (l_vecs, m_vecs), + "n_samples": n_samples, + "key": jax.random.PRNGKey(7), + "init_state_elems": elems, + "init_state_amps": amps, + "phase_fn": phase_fn, + } + # block_size 2 for weight-1 gates and 1 for heavier gates: scan path with padding + tiny = 2.5 * _qudit_phase_bytes_per_gate(1, n_samples, n_obs) / _BYTES_PER_GB + fn_single = build_qudit_expval_func(QuditCircuitConfig(max_memory=1.0, **common)) + fn_blocked = build_qudit_expval_func(QuditCircuitConfig(max_memory=tiny, **common)) + + params = jnp.array(rng.normal(size=n_gates)) + phase_params = jnp.array(rng.normal(size=n_qudits + 1)) + + for call_kwargs in [ + {}, + {"key": jax.random.PRNGKey(3), "n_samples": 128}, + {"observables": (l_vecs[:2], m_vecs[:2])}, + ]: + ev_s, cov_s = fn_single(params, phase_params, **call_kwargs) + ev_b, cov_b = fn_blocked(params, phase_params, **call_kwargs) + np.testing.assert_allclose(ev_b, ev_s, atol=1e-5, rtol=1e-5) + np.testing.assert_allclose(cov_b, cov_s, atol=1e-6, rtol=1e-4) + + def loss(fn): + return lambda p, q: jnp.real(jnp.sum(fn(p, q)[0])) + + grad_s = jax.grad(loss(fn_single), argnums=(0, 1))(params, phase_params) + grad_b = jax.grad(loss(fn_blocked), argnums=(0, 1))(params, phase_params) + for g_s, g_b in zip(grad_s, grad_b): + np.testing.assert_allclose(g_b, g_s, atol=1e-5, rtol=1e-5) + + def test_blocked_matches_exact(self): + """The blocked estimator still agrees with the brute-force reference within MC error.""" + d, n = 3, 3 + generators = np.array([[1, 0, 2], [0, 2, 1], [1, 1, 0], [2, 0, 0], [0, 1, 1]]) + thetas = np.array([0.4, -0.3, 0.7, 0.2, -0.5]) + l_vecs = np.array([[1, 0, 2], [0, 1, 0]]) + m_vecs = np.array([[0, 1, 0], [2, 0, 1]]) + gates = _dense_rows_to_sparse(generators) + tiny = 1.0 * _qudit_phase_bytes_per_gate(1, NUM_SAMPLES, 2) / _BYTES_PER_GB + config = QuditCircuitConfig( + dims=d, + n_qudits=n, + gates=gates, + observables=(l_vecs, m_vecs), + n_samples=NUM_SAMPLES, + key=jax.random.PRNGKey(5), + max_memory=tiny, + ) + mc_vals, mc_cov = build_qudit_expval_func(config)(jnp.array(thetas)) + for i, (l, m) in enumerate(zip(l_vecs, m_vecs)): + ref = qudit_expectation_brute_force(n, d, generators, thetas, l, m) + tol = max( + 3.5 * float(np.sqrt(mc_cov[i, 0, 0])), 3.5 * float(np.sqrt(mc_cov[i, 1, 1])), 1e-5 + ) + assert np.isclose(mc_vals[i], ref, atol=tol) From 9e21f741470b60dcf741634a81e8dd90ff531b4f Mon Sep 17 00:00:00 2001 From: Will Date: Thu, 8 Oct 2026 16:08:50 -0400 Subject: [PATCH 5/7] dtypes --- pennylane/labs/tcdq/mmd_loss_hw.py | 28 +++++--- pennylane/labs/tcdq/qudit_expval_functions.py | 72 ++++++++++--------- pennylane/labs/tcdq/utils.py | 3 +- 3 files changed, 57 insertions(+), 46 deletions(-) diff --git a/pennylane/labs/tcdq/mmd_loss_hw.py b/pennylane/labs/tcdq/mmd_loss_hw.py index 083e16879d8..c4d26f5ca3c 100644 --- a/pennylane/labs/tcdq/mmd_loss_hw.py +++ b/pennylane/labs/tcdq/mmd_loss_hw.py @@ -35,6 +35,10 @@ from .qudit_expval_functions import _dims_to_numpy +_REAL_DTYPE = jnp.float32 +_COMPLEX_DTYPE = jnp.complex64 +_INDEX_DTYPE = jnp.int32 + @dataclass(frozen=True) class QuditMMDConfig: @@ -93,7 +97,7 @@ def _cycle_marginal_probs(d: int, t: float) -> jnp.ndarray: For the derivation, see `Section IV B of Spectral Born machines: classically trainable quantum generative models for discrete data `_. """ - k = jnp.arange(d) + k = jnp.arange(d, dtype=_REAL_DTYPE) log_p = -4.0 * t * jnp.sin(jnp.pi * k / d) ** 2 p = jnp.exp(log_p) return p / jnp.sum(p) @@ -109,7 +113,7 @@ def _complete_marginal_probs(d: int, t: float) -> jnp.ndarray: For the derivation, see `Section IV B of Spectral Born machines: classically trainable quantum generative models for discrete data `_. """ - log_unnorm = jnp.zeros(d).at[1:].set(-t * d) + log_unnorm = jnp.zeros(d, dtype=_REAL_DTYPE).at[1:].set(-t * d) p = jnp.exp(log_unnorm) return p / jnp.sum(p) @@ -154,12 +158,14 @@ def _sample_fourier_indices( # pylint: disable=too-many-arguments,too-many-loca for wire in wire_tuple: wires_by_dim.setdefault(int(dims[wire]), []).append(int(wire)) - all_obs = jnp.zeros((n_ops, n_qudits), dtype=jnp.int32) + all_obs = jnp.zeros((n_ops, n_qudits), dtype=_INDEX_DTYPE) keys = jax.random.split(key, len(wires_by_dim)) if wires_by_dim else [] for group_key, (d_i, wires) in zip(keys, sorted(wires_by_dim.items())): marginal = _marginal_probs(d_i, bandwidth, graph_type) cols = jax.random.choice(group_key, d_i, shape=(n_ops, len(wires)), p=marginal) - all_obs = all_obs.at[:, jnp.asarray(wires, dtype=jnp.int32)].set(cols.astype(jnp.int32)) + all_obs = all_obs.at[:, jnp.asarray(wires, dtype=_INDEX_DTYPE)].set( + cols.astype(_INDEX_DTYPE) + ) return all_obs @@ -187,10 +193,10 @@ def _empirical_fourier_moments( Returns: Complex array of shape ``(n_obs,)``. """ - inv_d = 1.0 / jnp.asarray(dims_visible, dtype=float) - l_scaled = l_visible.astype(float) * inv_d[jnp.newaxis, :] - inner = l_scaled @ X_data.astype(float).T - return jnp.mean(jnp.exp(2j * jnp.pi * inner), axis=1) + inv_d = 1.0 / jnp.asarray(dims_visible, dtype=_REAL_DTYPE) + l_scaled = l_visible.astype(_REAL_DTYPE) * inv_d[jnp.newaxis, :] + inner = l_scaled @ X_data.astype(_REAL_DTYPE).T + return jnp.mean(jnp.exp(2j * jnp.pi * inner).astype(_COMPLEX_DTYPE), axis=1) def _pp_term(mu_p_hat: jnp.ndarray, m: int) -> jnp.ndarray: @@ -323,9 +329,9 @@ def _compute_qudit_loss_for_bandwidth( mu_q_hat, cov = model_output if isinstance(model_output, tuple) else (model_output, None) - mu_q_hat = jnp.asarray(mu_q_hat) + mu_q_hat = jnp.asarray(mu_q_hat).astype(_COMPLEX_DTYPE) if cov is not None: - cov = jnp.asarray(cov) + cov = jnp.asarray(cov).astype(_REAL_DTYPE) if mu_q_hat.shape != (n_ops,): raise ValueError( @@ -491,7 +497,7 @@ def loss_fn( ) active_key = jax.random.PRNGKey(0) if key is None else key - X_data = jnp.asarray(target_data) + X_data = jnp.asarray(target_data).astype(_INDEX_DTYPE) if X_data.ndim != 2: raise ValueError(f"target_data must be 2-D, got shape {X_data.shape}") diff --git a/pennylane/labs/tcdq/qudit_expval_functions.py b/pennylane/labs/tcdq/qudit_expval_functions.py index 99bd9ef39c5..48db1e5a4c1 100644 --- a/pennylane/labs/tcdq/qudit_expval_functions.py +++ b/pennylane/labs/tcdq/qudit_expval_functions.py @@ -38,6 +38,10 @@ import numpy as np from jax.typing import ArrayLike +_REAL_DTYPE = jnp.float32 +_COMPLEX_DTYPE = jnp.complex64 +_INDEX_DTYPE = jnp.int32 + @dataclass class QuditCircuitConfig: # pylint: disable=too-many-instance-attributes @@ -283,7 +287,7 @@ def _parse_qudit_gate_dict( omega=omega, supports=np.stack(sup_list), powers=np.stack(pow_list), - param_indices=jnp.array(pidx_list, dtype=int), + param_indices=jnp.array(pidx_list, dtype=_INDEX_DTYPE), ) for omega, (sup_list, pow_list, pidx_list) in sorted(buckets.items()) ] @@ -325,8 +329,11 @@ def _compute_qudit_samples( ) -> jnp.ndarray: """Generates uniformly random dit-strings from the product Z_{d_1} x ... x Z_{d_n}.""" - maxval = jnp.asarray(dims, dtype=jnp.int32)[jnp.newaxis, :] # (1, n_qudits) - return jax.random.randint(key, shape=(num_samples, n_qudits), minval=0, maxval=maxval) + maxval = jnp.asarray(dims, dtype=_INDEX_DTYPE)[jnp.newaxis, :] # (1, n_qudits) + # Draw with JAX's default integer dtype so the random stream for a given key is unchanged, + # then narrow to 32-bit for the downstream computation. + samples = jax.random.randint(key, shape=(num_samples, n_qudits), minval=0, maxval=maxval) + return samples.astype(_INDEX_DTYPE) _BYTES_PER_GB = 1024**3 @@ -361,12 +368,12 @@ def _obs_phase_matrix( :math:`J[i, j] = \\exp(i\\pi \\sum_k m_{ik} (2 z_{jk} - l_{ik}) / d_k)`. """ - s_f = samples.astype(jnp.float32) - inv_d = (1.0 / jnp.asarray(dims, dtype=jnp.float32))[jnp.newaxis, :] # (1, n_qudits) + s_f = samples.astype(_REAL_DTYPE) + inv_d = (1.0 / jnp.asarray(dims, dtype=_REAL_DTYPE))[jnp.newaxis, :] # (1, n_qudits) m_scaled = m_f * inv_d # (n_obs, n_qudits) return jnp.exp( 1j * jnp.pi * (2 * m_scaled @ s_f.T - jnp.sum(m_scaled * l_f, axis=1, keepdims=True)) - ) + ).astype(_COMPLEX_DTYPE) class _PrecomputedObsData(NamedTuple): @@ -413,7 +420,7 @@ def _group_block_contribution( omega = supports.shape[1] theta = theta_ext[param_indices] # (block,) - g = powers.astype(jnp.float32)[:, :, jnp.newaxis] # (block, omega, 1) + g = powers.astype(_REAL_DTYPE)[:, :, jnp.newaxis] # (block, omega, 1) d_s = dims_f[supports][:, :, jnp.newaxis] # (block, omega, 1) z_at_support = jnp.take(samples_t, supports, axis=0) # (block, omega, n_samples) l_at_support = jnp.take(l_t, supports, axis=0) # (block, omega, n_obs) @@ -452,9 +459,9 @@ def _group_phase_diffs( storing them, keeping gradient memory bounded as well. """ n_gates, omega = group.supports.shape - supports = jnp.asarray(group.supports, dtype=jnp.int32) - powers = jnp.asarray(group.powers, dtype=jnp.int32) - param_indices = jnp.asarray(group.param_indices, dtype=jnp.int32) + supports = jnp.asarray(group.supports, dtype=_INDEX_DTYPE) + powers = jnp.asarray(group.powers, dtype=_INDEX_DTYPE) + param_indices = jnp.asarray(group.param_indices, dtype=_INDEX_DTYPE) block_fn = jax.checkpoint( lambda s, p, i: _group_block_contribution(theta_ext, samples_t, l_t, dims_f, s, p, i) @@ -476,8 +483,7 @@ def _group_phase_diffs( def accumulate(total, block): return total + block_fn(*block), None - carry_dtype = jnp.result_type(theta_ext.dtype, samples_t.dtype) - zero = jnp.zeros((l_t.shape[1], samples_t.shape[1]), dtype=carry_dtype) + zero = jnp.zeros((l_t.shape[1], samples_t.shape[1]), dtype=_REAL_DTYPE) total, _ = jax.lax.scan( accumulate, zero, @@ -503,13 +509,13 @@ def _accumulate_phase_diffs( """Assemble the accumulated phase-difference matrix from all weight groups.""" n_samples = samples.shape[0] n_obs = l_vecs.shape[0] - theta = jnp.asarray(gates_params) - theta_ext = jnp.concatenate([theta, jnp.zeros((1,), dtype=theta.dtype)]) - samples_t = samples.astype(jnp.float32).T # (n_qudits, n_samples) - l_t = jnp.asarray(l_vecs).astype(jnp.float32).T # (n_qudits, n_obs) - dims_f = jnp.asarray(dims, dtype=jnp.float32) + theta = jnp.asarray(gates_params).astype(_REAL_DTYPE) + theta_ext = jnp.concatenate([theta, jnp.zeros((1,), dtype=_REAL_DTYPE)]) + samples_t = samples.astype(_REAL_DTYPE).T # (n_qudits, n_samples) + l_t = jnp.asarray(l_vecs).astype(_REAL_DTYPE).T # (n_qudits, n_obs) + dims_f = jnp.asarray(dims, dtype=_REAL_DTYPE) - accumulated = jnp.zeros((n_obs, n_samples)) + accumulated = jnp.zeros((n_obs, n_samples), dtype=_REAL_DTYPE) for group in gate_groups: block_size = _qudit_phase_block_size(max_memory_bytes, group.omega, n_samples, n_obs) accumulated = accumulated + _group_phase_diffs( @@ -517,7 +523,7 @@ def _accumulate_phase_diffs( ) if vmapped_phase_func is not None: - accumulated += vmapped_phase_func(phase_fn_params, samples, l_vecs) + accumulated += vmapped_phase_func(phase_fn_params, samples, l_vecs).astype(_REAL_DTYPE) return accumulated @@ -530,13 +536,13 @@ def _compute_initial_state_correction( dims: ArrayLike, ) -> jnp.ndarray: """Compute the correction factor for a non-standard initial state.""" - s_f = samples.astype(jnp.float32) - X_state = jnp.asarray(state_elems).astype(jnp.float32) # (N, n) - Psi = jnp.asarray(state_amps) # (N,) - inv_d = (1.0 / jnp.asarray(dims, dtype=jnp.float32))[jnp.newaxis, :] # (1, n) + s_f = samples.astype(_REAL_DTYPE) + X_state = jnp.asarray(state_elems).astype(_REAL_DTYPE) # (N, n) + Psi = jnp.asarray(state_amps).astype(_COMPLEX_DTYPE) # (N,) + inv_d = (1.0 / jnp.asarray(dims, dtype=_REAL_DTYPE))[jnp.newaxis, :] # (1, n) # ω^{Z·X^T} where ω_j = exp(2πi/d_j) — shape (s, N) - omega_ZX = jnp.exp(2j * jnp.pi * ((s_f * inv_d) @ X_state.T)) + omega_ZX = jnp.exp(2j * jnp.pi * ((s_f * inv_d) @ X_state.T)).astype(_COMPLEX_DTYPE) # Ψ̃^(2) = ω^{Z·X^T} · Ψ — shape (s,) psi_tilde_2 = omega_ZX @ Psi @@ -545,7 +551,7 @@ def _compute_initial_state_correction( F_mat = Psi.conj()[:, jnp.newaxis] * omega_ZX.conj().T # Ψ̃^(1) = ω^{L·X^T} · F — shape (l, s) - omega_LX = jnp.exp(2j * jnp.pi * ((l_f * inv_d) @ X_state.T)) # (l, N) + omega_LX = jnp.exp(2j * jnp.pi * ((l_f * inv_d) @ X_state.T)).astype(_COMPLEX_DTYPE) # (l, N) psi_tilde_1 = omega_LX @ F_mat # H = Ψ̃^(1) ⊙ (1_{l×1} · (Ψ̃^(2))^T) — shape (l, s) @@ -669,10 +675,10 @@ def compute_phase_diff(p_params, sample, l_vec): ) if config.observables is not None: - l_vecs = jnp.array(config.observables[0], dtype=jnp.int32) - m_vecs = jnp.array(config.observables[1], dtype=jnp.int32) - l_f = l_vecs.astype(jnp.float32) - m_f = m_vecs.astype(jnp.float32) + l_vecs = jnp.array(config.observables[0], dtype=_INDEX_DTYPE) + m_vecs = jnp.array(config.observables[1], dtype=_INDEX_DTYPE) + l_f = l_vecs.astype(_REAL_DTYPE) + m_f = m_vecs.astype(_REAL_DTYPE) n_obs = l_vecs.shape[0] defaults = _PrecomputedObsData( l_vecs=l_vecs, @@ -720,9 +726,9 @@ def qudit_expval_batched( of the mean estimator, shape ``(n_obs, 2, 2)``. """ if observables is not None: - l_vecs = jnp.array(observables[0], dtype=jnp.int32) - l_f = l_vecs.astype(jnp.float32) - m_f = jnp.array(observables[1], dtype=jnp.int32).astype(jnp.float32) + l_vecs = jnp.array(observables[0], dtype=_INDEX_DTYPE) + l_f = l_vecs.astype(_REAL_DTYPE) + m_f = jnp.array(observables[1], dtype=_INDEX_DTYPE).astype(_REAL_DTYPE) elif defaults is not None: l_vecs, l_f, m_f = defaults.l_vecs, defaults.l_f, defaults.m_f else: @@ -761,7 +767,7 @@ def qudit_expval_batched( state_elems = config.init_state_elems if init_state_elems is None else init_state_elems state_amps = config.init_state_amps if init_state_amps is None else init_state_amps - integrand = obs_pm * jnp.exp(1j * accumulated_phase_diffs) + integrand = obs_pm * jnp.exp(1j * accumulated_phase_diffs).astype(_COMPLEX_DTYPE) if state_elems is not None and state_amps is not None: H = _compute_initial_state_correction(samples, l_f, state_elems, state_amps, dims) integrand = integrand * H diff --git a/pennylane/labs/tcdq/utils.py b/pennylane/labs/tcdq/utils.py index 405133b9e1a..b9e3188268a 100644 --- a/pennylane/labs/tcdq/utils.py +++ b/pennylane/labs/tcdq/utils.py @@ -260,8 +260,7 @@ def create_random_qudit_gates( rng = np.random.default_rng(seed) gates_dict = {} - from tqdm import tqdm - for i in tqdm(range(n_gates)): + for i in range(n_gates): weight = int(rng.integers(min_weight, max_weight + 1)) support = sorted(int(q) for q in rng.choice(n_qudits, size=weight, replace=False)) gates_dict[i] = [{q: int(rng.integers(1, dims[q])) for q in support}] From adfacc1d5ad2a5deb913ad7adfd2f5ff0be53a38 Mon Sep 17 00:00:00 2001 From: Will Date: Tue, 20 Oct 2026 15:37:44 -0400 Subject: [PATCH 6/7] fix merge --- pennylane/labs/tcdq/qudit_expval_functions.py | 20 +++++++++++++++---- pennylane/labs/tests/tcdq/test_mmd_loss_hw.py | 1 - 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/pennylane/labs/tcdq/qudit_expval_functions.py b/pennylane/labs/tcdq/qudit_expval_functions.py index 0c8960a5819..ebfecf62507 100644 --- a/pennylane/labs/tcdq/qudit_expval_functions.py +++ b/pennylane/labs/tcdq/qudit_expval_functions.py @@ -292,6 +292,7 @@ def _compute_qudit_samples( samples = jax.random.randint(key, shape=(num_samples, n_qudits), minval=0, maxval=maxval) return samples.astype(_INDEX_DTYPE) + _BYTES_PER_GB = 1024**3 _FLOAT_BYTES = 4 @@ -472,10 +473,21 @@ def _accumulate_phase_diffs( accumulated = jnp.zeros((n_obs, n_samples), dtype=_REAL_DTYPE) for group in gate_groups: - block_size = _qudit_phase_block_size(max_memory_bytes, group.omega, n_samples, n_obs) - accumulated = accumulated + _group_phase_diffs( - theta_ext, samples_t, l_t, dims_f, group, block_size - ) + if max_memory_bytes is None: + accumulated = accumulated + _group_block_contribution( + theta_ext, + samples_t, + l_t, + dims_f, + jnp.asarray(group.supports, dtype=_INDEX_DTYPE), + jnp.asarray(group.powers, dtype=_INDEX_DTYPE), + jnp.asarray(group.param_indices, dtype=_INDEX_DTYPE), + ) + else: + block_size = _qudit_phase_block_size(max_memory_bytes, group.omega, n_samples, n_obs) + accumulated = accumulated + _group_phase_diffs( + theta_ext, samples_t, l_t, dims_f, group, block_size + ) if vmapped_phase_func is not None: accumulated += vmapped_phase_func(phase_fn_params, samples, l_vecs).astype(_REAL_DTYPE) diff --git a/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py b/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py index b67c048210e..3bf22876f84 100644 --- a/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py +++ b/pennylane/labs/tests/tcdq/test_mmd_loss_hw.py @@ -31,7 +31,6 @@ from pennylane.labs.tcdq.qudit_expval_functions import ( QuditCircuitConfig, _dims_to_numpy, - _gates_to_dense_generators, build_qudit_expval_func, ) From 9debd024424310212c75ab389f9395522ce4eb15 Mon Sep 17 00:00:00 2001 From: Will Date: Tue, 20 Oct 2026 15:38:56 -0400 Subject: [PATCH 7/7] remove comment --- pennylane/labs/tcdq/qudit_expval_functions.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pennylane/labs/tcdq/qudit_expval_functions.py b/pennylane/labs/tcdq/qudit_expval_functions.py index ebfecf62507..134830e6bb4 100644 --- a/pennylane/labs/tcdq/qudit_expval_functions.py +++ b/pennylane/labs/tcdq/qudit_expval_functions.py @@ -287,8 +287,6 @@ def _compute_qudit_samples( """Generates uniformly random dit-strings from the product Z_{d_1} x ... x Z_{d_n}.""" maxval = jnp.asarray(dims, dtype=_INDEX_DTYPE)[jnp.newaxis, :] # (1, n_qudits) - # Draw with JAX's default integer dtype so the random stream for a given key is unchanged, - # then narrow to 32-bit for the downstream computation. samples = jax.random.randint(key, shape=(num_samples, n_qudits), minval=0, maxval=maxval) return samples.astype(_INDEX_DTYPE)