[Labs] TCDQ qudit update dtypes

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 <https://arxiv.org/abs/2607.06675>`_.
     """
-    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 <https://arxiv.org/abs/2607.06675>`_.
     """
-    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)
 
@@ -150,12 +154,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
 
 
@@ -183,10 +189,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:
@@ -319,9 +325,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(
@@ -487,7 +493,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}")

pennylane/labs/tcdq/qudit_expval_functions.py

@@ -40,6 +40,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
@@ -271,7 +275,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())
     ]
@@ -282,8 +286,9 @@ 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)
+    samples = jax.random.randint(key, shape=(num_samples, n_qudits), minval=0, maxval=maxval)
+    return samples.astype(_INDEX_DTYPE)
 
 
 _BYTES_PER_GB = 1024**3
@@ -326,12 +331,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):
@@ -372,7 +377,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)
@@ -407,9 +412,9 @@ def _group_phase_diffs(
 ) -> jnp.ndarray:
     """Phase differences of one weight group, processed in memory-bounded blocks."""
     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)
@@ -431,8 +436,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,
@@ -459,23 +463,23 @@ 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:
         if max_memory_bytes is None:
             accumulated = accumulated + _group_block_contribution(
                 theta_ext,
                 samples_t,
                 l_t,
                 dims_f,
-                jnp.asarray(group.supports, dtype=jnp.int32),
-                jnp.asarray(group.powers, dtype=jnp.int32),
-                jnp.asarray(group.param_indices, dtype=jnp.int32),
+                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)
@@ -484,7 +488,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
 
@@ -497,13 +501,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
@@ -512,7 +516,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)
@@ -636,10 +640,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,
@@ -687,9 +691,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:
@@ -728,7 +732,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