From 0a9c1e92fdc59b519fb825dd8bda37d57875ca38 Mon Sep 17 00:00:00 2001
From: GetlinDu <1578440935@qq.com>
Date: Thu, 9 Jul 2026 19:04:18 +0800
Subject: [PATCH 1/4] Add Kensingtonian function

---
 src/deepquantum/photonic/kensingtonian_.py | 122 +++++++++++++++++++++
 1 file changed, 122 insertions(+)
 create mode 100644 src/deepquantum/photonic/kensingtonian_.py

diff --git a/src/deepquantum/photonic/kensingtonian_.py b/src/deepquantum/photonic/kensingtonian_.py
new file mode 100644
index 00000000..0d0edf03
--- /dev/null
+++ b/src/deepquantum/photonic/kensingtonian_.py
@@ -0,0 +1,122 @@
+"""Functions for Kensingtonian"""
+
+from itertools import product
+
+import torch
+
+
+def _multinomial_click_coeff(num_detectors: int, click: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
+    """Return N! / ((N-k)! (k-d)! d!) for each mode."""
+    n = torch.as_tensor(num_detectors, dtype=click.dtype, device=click.device)
+    return torch.exp(
+        torch.lgamma(n + 1) - torch.lgamma(n - click + 1) - torch.lgamma(click - d + 1) - torch.lgamma(d + 1)
+    )
+
+
+def _mode_submat(mat: torch.Tensor, keep_modes: torch.Tensor) -> torch.Tensor:
+    """Keep both quadratures of each selected mode."""
+    if keep_modes.numel() == 0:
+        return mat.new_empty((0, 0))
+    idx1 = keep_modes
+    idx2 = idx1 + mat.shape[-1] // 2
+    idx = torch.sort(torch.cat([idx1, idx2]))[0]
+    return mat[idx][:, idx]
+
+
+def _mode_subvec(vec: torch.Tensor, keep_modes: torch.Tensor) -> torch.Tensor:
+    """Keep both quadratures of each selected mode from a vector."""
+    if keep_modes.numel() == 0:
+        return vec.new_empty((0,))
+    idx1 = keep_modes
+    idx2 = idx1 + vec.shape[-1] // 2
+    idx = torch.sort(torch.cat([idx1, idx2]))[0]
+    return vec[idx]
+
+
+def _dz_diag(d: torch.Tensor, keep_modes: torch.Tensor, num_detectors: int) -> torch.Tensor:
+    """Construct D_Z from Eq. (28) on the retained modes."""
+    values = (num_detectors - d[keep_modes]) / d[keep_modes]
+    return torch.cat([values, values]).diag()
+
+
+def _kensingtonian_term(
+    matrix: torch.Tensor,
+    clicks: torch.Tensor,
+    d: torch.Tensor,
+    num_detectors: int,
+    alpha: torch.Tensor | None,
+) -> torch.Tensor:
+    sign_power = torch.sum(clicks - d).to(torch.long)
+    sign = -1 if int(sign_power.item()) % 2 else 1
+    coeff = _multinomial_click_coeff(num_detectors, clicks, d).prod()
+    keep_modes = torch.nonzero(d > 0, as_tuple=False).flatten()
+    if keep_modes.numel() == 0:
+        det_term = matrix.new_tensor(1)
+        exp_term = matrix.new_tensor(1)
+    else:
+        size = matrix.shape[-1]
+        identity = torch.eye(size, dtype=matrix.dtype, device=matrix.device)
+        base = _mode_submat(identity - matrix, keep_modes)
+        dz = _dz_diag(d, keep_modes, num_detectors).to(matrix.dtype)
+        det_mat = base + dz
+        det_term = torch.linalg.det(det_mat)
+        if alpha is None:
+            exp_term = matrix.new_tensor(1)
+        else:
+            rhs = _mode_subvec((identity - matrix) @ alpha, keep_modes)
+            quad = rhs @ torch.linalg.solve(det_mat, rhs)
+            exp_term = torch.exp(quad)
+    prefactor = (num_detectors / d[keep_modes]).prod() if keep_modes.numel() > 0 else matrix.new_tensor(1)
+    return sign * coeff.to(matrix.dtype) * prefactor.to(matrix.dtype) * exp_term / torch.sqrt(det_term)
+
+
+def kensingtonian(
+    matrix: torch.Tensor,
+    clicks: torch.Tensor,
+    num_detectors: int,
+    alpha: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Calculate the Kensingtonian or loop Kensingtonian"""
+    assert matrix.dim() == 2
+    assert matrix.shape[-2] == matrix.shape[-1]
+    assert matrix.shape[-1] % 2 == 0
+    assert num_detectors >= 1
+    m = matrix.shape[-1] // 2
+    clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device)
+    assert clicks.shape == (m,)
+    assert torch.all(clicks >= 0)
+    assert torch.all(clicks <= num_detectors)
+    clicks_float = clicks.to(matrix.real.dtype)
+    if alpha is not None:
+        assert alpha.shape == (2 * m,)
+        alpha = alpha.to(dtype=matrix.dtype, device=matrix.device)
+    ken = matrix.new_tensor(0)
+    ranges = [range(int(click.item()) + 1) for click in clicks]
+    for d_tuple in product(*ranges):
+        d = torch.tensor(d_tuple, dtype=clicks_float.dtype, device=matrix.device)
+        ken = ken + _kensingtonian_term(matrix, clicks_float, d, num_detectors, alpha)
+    return ken
+
+
+def kensingtonian_batch(
+    matrix: torch.Tensor,
+    clicks: torch.Tensor,
+    num_detectors: int,
+    alpha: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Calculate batched Kensingtonians."""
+    assert matrix.dim() == 3
+    assert matrix.shape[-2] == matrix.shape[-1]
+    assert matrix.shape[-1] % 2 == 0
+    clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device)
+    if clicks.dim() == 1:
+        clicks = clicks.expand(matrix.shape[0], -1)
+    if alpha is not None:
+        alpha = torch.as_tensor(alpha, dtype=matrix.dtype, device=matrix.device)
+        if alpha.dim() == 1:
+            alpha = alpha.expand(matrix.shape[0], -1)
+    values = []
+    for i in range(matrix.shape[0]):
+        alpha_i = None if alpha is None else alpha[i]
+        values.append(kensingtonian(matrix[i], clicks[i], num_detectors, alpha_i))
+    return torch.stack(values)

From 9644a5086823dfdf77fb315d95f131fa9bccb692 Mon Sep 17 00:00:00 2001
From: GetlinDu <1578440935@qq.com>
Date: Fri, 10 Jul 2026 18:08:33 +0800
Subject: [PATCH 2/4] Optimize Kensingtonian calculation

---
 src/deepquantum/photonic/kensingtonian_.py | 37 ++++++++++++++++------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/deepquantum/photonic/kensingtonian_.py b/src/deepquantum/photonic/kensingtonian_.py
index 0d0edf03..0c682227 100644
--- a/src/deepquantum/photonic/kensingtonian_.py
+++ b/src/deepquantum/photonic/kensingtonian_.py
@@ -1,10 +1,17 @@
 """Functions for Kensingtonian"""
 
+from functools import cache
 from itertools import product
 
 import torch
 
 
+@cache
+def _d_tuples(clicks: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+    ranges = [range(click + 1) for click in clicks]
+    return tuple(product(*ranges))
+
+
 def _multinomial_click_coeff(num_detectors: int, click: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
     """Return N! / ((N-k)! (k-d)! d!) for each mode."""
     n = torch.as_tensor(num_detectors, dtype=click.dtype, device=click.device)
@@ -41,31 +48,33 @@ def _dz_diag(d: torch.Tensor, keep_modes: torch.Tensor, num_detectors: int) -> t
 
 def _kensingtonian_term(
     matrix: torch.Tensor,
+    base_full: torch.Tensor,
     clicks: torch.Tensor,
     d: torch.Tensor,
     num_detectors: int,
-    alpha: torch.Tensor | None,
+    beta: torch.Tensor | None,
 ) -> torch.Tensor:
     sign_power = torch.sum(clicks - d).to(torch.long)
     sign = -1 if int(sign_power.item()) % 2 else 1
+
     coeff = _multinomial_click_coeff(num_detectors, clicks, d).prod()
     keep_modes = torch.nonzero(d > 0, as_tuple=False).flatten()
     if keep_modes.numel() == 0:
         det_term = matrix.new_tensor(1)
         exp_term = matrix.new_tensor(1)
     else:
-        size = matrix.shape[-1]
-        identity = torch.eye(size, dtype=matrix.dtype, device=matrix.device)
-        base = _mode_submat(identity - matrix, keep_modes)
+        base = _mode_submat(base_full, keep_modes)
         dz = _dz_diag(d, keep_modes, num_detectors).to(matrix.dtype)
         det_mat = base + dz
         det_term = torch.linalg.det(det_mat)
-        if alpha is None:
+
+        if beta is None:
             exp_term = matrix.new_tensor(1)
         else:
-            rhs = _mode_subvec((identity - matrix) @ alpha, keep_modes)
+            rhs = _mode_subvec(beta, keep_modes)
             quad = rhs @ torch.linalg.solve(det_mat, rhs)
             exp_term = torch.exp(quad)
+
     prefactor = (num_detectors / d[keep_modes]).prod() if keep_modes.numel() > 0 else matrix.new_tensor(1)
     return sign * coeff.to(matrix.dtype) * prefactor.to(matrix.dtype) * exp_term / torch.sqrt(det_term)
 
@@ -76,25 +85,31 @@ def kensingtonian(
     num_detectors: int,
     alpha: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Calculate the Kensingtonian or loop Kensingtonian"""
+    """Calculate the Kensingtonian or loop Kensingtonian."""
     assert matrix.dim() == 2
     assert matrix.shape[-2] == matrix.shape[-1]
     assert matrix.shape[-1] % 2 == 0
     assert num_detectors >= 1
+
     m = matrix.shape[-1] // 2
     clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device)
     assert clicks.shape == (m,)
     assert torch.all(clicks >= 0)
     assert torch.all(clicks <= num_detectors)
+
     clicks_float = clicks.to(matrix.real.dtype)
     if alpha is not None:
         assert alpha.shape == (2 * m,)
         alpha = alpha.to(dtype=matrix.dtype, device=matrix.device)
+
+    identity = torch.eye(matrix.shape[-1], dtype=matrix.dtype, device=matrix.device)
+    base_full = identity - matrix
+    beta = None if alpha is None else base_full @ alpha
     ken = matrix.new_tensor(0)
-    ranges = [range(int(click.item()) + 1) for click in clicks]
-    for d_tuple in product(*ranges):
+    clicks_tuple = tuple(int(click.item()) for click in clicks)
+    for d_tuple in _d_tuples(clicks_tuple):
         d = torch.tensor(d_tuple, dtype=clicks_float.dtype, device=matrix.device)
-        ken = ken + _kensingtonian_term(matrix, clicks_float, d, num_detectors, alpha)
+        ken = ken + _kensingtonian_term(matrix, base_full, clicks_float, d, num_detectors, beta)
     return ken
 
 
@@ -108,6 +123,7 @@ def kensingtonian_batch(
     assert matrix.dim() == 3
     assert matrix.shape[-2] == matrix.shape[-1]
     assert matrix.shape[-1] % 2 == 0
+
     clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device)
     if clicks.dim() == 1:
         clicks = clicks.expand(matrix.shape[0], -1)
@@ -115,6 +131,7 @@ def kensingtonian_batch(
         alpha = torch.as_tensor(alpha, dtype=matrix.dtype, device=matrix.device)
         if alpha.dim() == 1:
             alpha = alpha.expand(matrix.shape[0], -1)
+
     values = []
     for i in range(matrix.shape[0]):
         alpha_i = None if alpha is None else alpha[i]

From e06de98011f26197ba82742de8573e6708fd13ea Mon Sep 17 00:00:00 2001
From: GetlinDu <1578440935@qq.com>
Date: Tue, 14 Jul 2026 17:43:27 +0800
Subject: [PATCH 3/4] kensingtonian modified

---
 src/deepquantum/photonic/__init__.py       |   2 +
 src/deepquantum/photonic/circuit.py        | 134 +++++++++++---
 src/deepquantum/photonic/kensingtonian_.py | 202 ++++++++++++++-------
 3 files changed, 248 insertions(+), 90 deletions(-)

diff --git a/src/deepquantum/photonic/__init__.py b/src/deepquantum/photonic/__init__.py
index 69f95193..9f5a52f3 100644
--- a/src/deepquantum/photonic/__init__.py
+++ b/src/deepquantum/photonic/__init__.py
@@ -9,6 +9,7 @@
     draw,
     gate,
     hafnian_,
+    kensingtonian_,
     mapper,
     measurement,
     operation,
@@ -47,6 +48,7 @@
     UAnyGate,
 )
 from .hafnian_ import hafnian
+from .kensingtonian_ import kensingtonian
 from .mapper import UnitaryMapper
 from .measurement import GeneralBosonic, Generaldyne, Homodyne, PhotonNumberResolvingBosonic
 from .qmath import (
diff --git a/src/deepquantum/photonic/circuit.py b/src/deepquantum/photonic/circuit.py
index 23788eb1..2acd5c74 100644
--- a/src/deepquantum/photonic/circuit.py
+++ b/src/deepquantum/photonic/circuit.py
@@ -43,6 +43,7 @@
     UAnyGate,
 )
 from .hafnian_ import hafnian
+from .kensingtonian_ import kensingtonian
 from .measurement import Generaldyne, Homodyne
 from .operation import Channel, Delay, Gate, Operation
 from .qmath import (
@@ -1077,10 +1078,14 @@ def _get_prob_fock_vmap(self, sub_mat: torch.Tensor, per_norm: torch.Tensor) ->
         prob = torch.abs(amplitude) ** 2
         return prob
 
-    def _get_prob_gaussian(self, final_state: Any, state: Any = None) -> torch.Tensor:
+    def _get_prob_gaussian(
+        self, final_state: Any, state: Any = None, detector: str | None = None, num_detectors: int = 1
+    ) -> torch.Tensor:
         """Get the batched probabilities of the final state for Gaussian backend."""
         if not isinstance(final_state, torch.Tensor):
             final_state = torch.tensor(final_state, dtype=torch.long)
+            detector = self.detector if detector is None else detector
+            num_detectors = getattr(self, '_num_detectors', num_detectors)
         if state is None:
             cov = self._cov
             mean = self._mean
@@ -1097,7 +1102,9 @@ def _get_prob_gaussian(self, final_state: Any, state: Any = None) -> torch.Tenso
         batch = cov.shape[0]
         probs = []
         for i in range(batch):
-            prob = self._get_probs_gaussian_helper(final_state, cov=cov[i], mean=mean[i], detector=self.detector)[0]
+            prob = self._get_probs_gaussian_helper(
+                final_state, cov=cov[i], mean=mean[i], detector=detector, num_detectors=num_detectors
+            )[0]
             probs.append(prob)
         return torch.stack(probs).squeeze()
 
@@ -1109,6 +1116,7 @@ def _get_probs_gaussian_helper(
         detector: str = 'pnrd',
         purity: bool | None = None,
         loop: bool | None = None,
+        num_detectors: int = 1,
     ) -> torch.Tensor:
         """Get the probabilities of the final states for Gaussian backend."""
         if loop is None:
@@ -1129,13 +1137,32 @@ def _get_probs_gaussian_helper(
         gamma = mean_ladder.mH @ q.inverse()
         if detector == 'pnrd':
             matrix = a_mat
-        elif detector == 'threshold':
+        elif detector in ('threshold', 'click_counting'):
             matrix = o_mat
+        else:
+            raise ValueError(f'Unsupported detector: {detector}')
         if purity is None:
             purity = GaussianState([cov, mean]).is_pure
         p_vac = torch.exp(-0.5 * mean_ladder.mH @ q.inverse() @ mean_ladder) / det_q.sqrt()
-        batch_get_prob = vmap(self._get_prob_gaussian_base, in_dims=(0, None, None, None, None, None, None))
-        probs = batch_get_prob(final_states, matrix, gamma, p_vac, detector, purity, loop)
+        if detector == 'click_counting':
+            probs = torch.stack(
+                [
+                    self._get_prob_gaussian_base(
+                        final_state,
+                        matrix,
+                        gamma,
+                        p_vac,
+                        detector,
+                        purity,
+                        loop,
+                        num_detectors,
+                    )
+                    for final_state in final_states
+                ]
+            )
+        else:
+            batch_get_prob = vmap(self._get_prob_gaussian_base, in_dims=(0, None, None, None, None, None, None, None))
+            probs = batch_get_prob(final_states, matrix, gamma, p_vac, detector, purity, loop, num_detectors)
         return probs
 
     def _get_prob_gaussian_base(
@@ -1147,6 +1174,7 @@ def _get_prob_gaussian_base(
         detector: str = 'pnrd',
         purity: bool = True,
         loop: bool = False,
+        num_detectors: int = 1,
     ) -> torch.Tensor:
         """Get the probability of the final state for Gaussian backend."""
         gamma = gamma.squeeze()
@@ -1175,6 +1203,9 @@ def _get_prob_gaussian_base(
             final_state_double = torch.cat([final_state, final_state])
             sub_mat = sub_matrix(matrix, final_state_double, final_state_double)
             prob = p_vac * torontonian(sub_mat, sub_gamma)
+        elif detector == 'click_counting':
+            alpha = gamma if loop else None
+            prob = p_vac * kensingtonian(matrix, final_state, num_detectors, alpha)
         return abs(prob.real).squeeze()
 
     def _get_prob_mps(self, final_state: Any, wires: int | list[int] | None = None) -> torch.Tensor:
@@ -1203,6 +1234,7 @@ def measure(
         wires: int | list[int] | None = None,
         detector: str | None = None,
         mcmc: bool = False,
+        num_detectors: int = 1,
     ) -> dict | list[dict] | None:
         """Measure the final state.
 
@@ -1215,6 +1247,7 @@ def measure(
             detector: For Gaussian backend, use ``'pnrd'`` for the photon-number-resolving detector or
                 ``'threshold'`` for the threshold detector. Default: ``None``
             mcmc: Whether to use MCMC sampling method. Default: ``False``
+            num_detectors: The number of threshold detectors in each click-counting detector. Default: 1
 
         See https://arxiv.org/pdf/2108.01622 for MCMC.
         """
@@ -1228,7 +1261,7 @@ def measure(
             results = self._measure_fock(shots, with_prob, wires, mcmc)
         elif self.backend == 'gaussian':
             detector = self.detector if detector is None else detector.lower()
-            results = self._measure_gaussian(shots, with_prob, wires, detector, mcmc)
+            results = self._measure_gaussian(shots, with_prob, wires, detector, mcmc, num_detectors)
         if len(results) == 1:
             results = results[0]
         return results
@@ -1428,10 +1461,18 @@ def _sample_mcmc_fock(self, shots: int, init_state: torch.Tensor, unitary: torch
         )
         return merged_samples
 
-    def _measure_gaussian(self, shots: int, with_prob: bool, wires: list[int], detector: str, mcmc: bool) -> list[dict]:
+    def _measure_gaussian(
+        self,
+        shots: int,
+        with_prob: bool,
+        wires: list[int],
+        detector: str,
+        mcmc: bool,
+        num_detectors: int = 1,
+    ) -> list[dict]:
         """Measure the final state for Gaussian backend."""
         if isinstance(self.state, list):
-            return self._measure_gaussian_state(shots, with_prob, wires, detector, mcmc)
+            return self._measure_gaussian_state(shots, with_prob, wires, detector, mcmc, num_detectors)
         elif isinstance(self.state, dict):
             assert not mcmc, "Final states have been calculated, we don't need mcmc!"
             print('Automatically using the default detector!')
@@ -1440,7 +1481,13 @@ def _measure_gaussian(self, shots: int, with_prob: bool, wires: list[int], detec
             raise ValueError('Check your forward function or input!')
 
     def _measure_gaussian_state(
-        self, shots: int, with_prob: bool, wires: list[int], detector: str, mcmc: bool
+        self,
+        shots: int,
+        with_prob: bool,
+        wires: list[int],
+        detector: str,
+        mcmc: bool,
+        num_detectors: int = 1,
     ) -> list[dict]:
         """Measure the final state according to Gaussian state for Gaussian backend.
 
@@ -1455,14 +1502,19 @@ def _measure_gaussian_state(
             print('Using MCMC method to sample the final states!')
             for i in range(batch):
                 samples_i = self._sample_mcmc_gaussian(
-                    shots=shots, cov=cov[i], mean=mean[i], detector=detector, num_chain=5
+                    shots=shots,
+                    cov=cov[i],
+                    mean=mean[i],
+                    detector=detector,
+                    num_chain=5,
+                    num_detectors=num_detectors,
                 )
                 all_samples.append(samples_i)
         else:  # chain-rule method with small number of shots
             print('Using chain-rule method to sample the final states!')
             samples = []
             for _ in range(shots):
-                sample = self._generate_chain_sample(wires, detector)
+                sample = self._generate_chain_sample(wires, detector, num_detectors)
                 samples.append(sample)
             samples = torch.stack(samples).permute(1, 0, 2)  # (batch, shots, wires)
             for i in range(batch):
@@ -1475,11 +1527,16 @@ def _measure_gaussian_state(
             if with_prob:
                 for k in samples_i:
                     if mcmc:
-                        prob = self._get_prob_gaussian(k, [cov[i], mean[i]])
+                        prob = self._get_prob_gaussian(k, [cov[i], mean[i]], num_detectors=num_detectors)
                     else:
                         wires_ = torch.tensor(wires, device=cov.device)
                         idx = torch.cat([wires_, wires_ + self.nmode])
-                        prob = self._get_prob_gaussian(k, [cov[i][idx[:, None], idx], mean[i][idx, :]])
+                        prob = self._get_prob_gaussian(
+                            k,
+                            [cov[i][idx[:, None], idx], mean[i][idx, :]],
+                            detector=detector,
+                            num_detectors=num_detectors,
+                        )
                     samples_i[k] = samples_i[k], prob
             for key in samples_i:
                 state_b = [key[wire] for wire in wires] if mcmc else list(key)
@@ -1495,11 +1552,20 @@ def _measure_gaussian_state(
             all_results.append(results)
         return all_results
 
-    def _sample_mcmc_gaussian(self, shots: int, cov: torch.Tensor, mean: torch.Tensor, detector: str, num_chain: int):
+    def _sample_mcmc_gaussian(
+        self,
+        shots: int,
+        cov: torch.Tensor,
+        mean: torch.Tensor,
+        detector: str,
+        num_chain: int,
+        num_detectors: int = 1,
+    ):
         """Sample the output states for Gaussian backend via SC-MCMC method."""
         self._cov = cov
         self._mean = mean
         self.detector = detector
+        self._num_detectors = num_detectors
         if detector == 'threshold' and not torch.allclose(mean, torch.zeros_like(mean)):
             # For the displaced state, aggregate PNRD detector samples to derive threshold detector results
             self.detector = 'pnrd'
@@ -1530,25 +1596,31 @@ def _proposal_sampler(self):
             assert self.basis, 'Currently NOT supported.'
             sample = self._out_fock_basis[torch.randint(0, len(self._out_fock_basis), (1,))[0]]
         elif self.backend == 'gaussian':
-            sample = self._generate_rand_sample(self.detector)
+            num_detectors = getattr(self, '_num_detectors', 1)
+            sample = self._generate_rand_sample(self.detector, num_detectors)
         return tuple(sample.tolist())
 
-    def _generate_rand_sample(self, detector: str = 'pnrd'):
+    def _generate_rand_sample(self, detector: str = 'pnrd', num_detectors: int = 1):
         """Generate a random sample according to uniform proposal distribution."""
         nmode = self._nmode_tdm if self._with_delay else self.nmode
         if detector == 'threshold':
             sample = torch.randint(0, 2, [nmode])
+        elif detector == 'click_counting':
+            sample = torch.randint(0, num_detectors + 1, [nmode])
         elif detector == 'pnrd':
             sample = torch.randint(0, self.cutoff, [nmode])
+        else:
+            raise ValueError(f'Unsupported detector: {detector}')
         return sample
 
-    def _generate_chain_sample(self, wires: list[int], detector: str) -> torch.Tensor:
+    def _generate_chain_sample(self, wires: list[int], detector: str, num_detectors: int = 1) -> torch.Tensor:
         """Generate batched random samples via chain rule.
 
         Args:
             wires: The wires to measure. It can be a list of integers specifying the indices of the wires.
-            detector: For Gaussian backend, use ``'pnrd'`` for the photon-number-resolving detector or
-                ``'threshold'`` for the threshold detector.
+            detector: For Gaussian backend, use ``'pnrd'`` for the photon-number-resolving detector,
+                ``'threshold'`` for the threshold detector, or ``'click_counting'`` for the click-counting detector.
+            num_detectors: The number of threshold detectors in each click-counting detector. Default: 1
 
         Returns:
             Tensor of shape (batch, nwire).
@@ -1567,10 +1639,10 @@ def _generate_chain_sample(self, wires: list[int], detector: str) -> torch.Tenso
                 mps[i] = torch.gather(mps[i], dim=2, index=index)
             sample = torch.stack(sample, dim=-1).squeeze(1)
         elif self.backend == 'gaussian':  # chain rule for GBS
-            sample = self._generate_chain_sample_gaussian(wires, detector)
+            sample = self._generate_chain_sample_gaussian(wires, detector, num_detectors)
         return sample
 
-    def _generate_chain_sample_gaussian(self, wires: list[int], detector: str) -> torch.Tensor:
+    def _generate_chain_sample_gaussian(self, wires: list[int], detector: str, num_detectors: int = 1) -> torch.Tensor:
         """Generate batched random samples via chain rule for Gaussian backend.
 
         See
@@ -1581,7 +1653,16 @@ def _generate_chain_sample_gaussian(self, wires: list[int], detector: str) -> to
         def _sample_wire(sample, cov_sub, mean_sub, cutoff, detector):
             """Sample for a wire."""
             states = [torch.tensor(sample + [i], device=cov_sub.device) for i in range(cutoff)]
-            probs = [self._get_probs_gaussian_helper(s, cov_sub, mean_sub, detector) for s in states]
+            probs = [
+                self._get_probs_gaussian_helper(
+                    s,
+                    cov_sub,
+                    mean_sub,
+                    detector,
+                    num_detectors=num_detectors,
+                )
+                for s in states
+            ]
             sample_wire = torch.multinomial(torch.cat(probs), num_samples=1)
             return sample_wire
 
@@ -1633,7 +1714,14 @@ def _sample_mixed(cov, mean, wires, nmode, cutoff, detector, eps=5e-5):
         purity = GaussianState(self.state).is_pure
         cov, mean = self.state
         batch = cov.shape[0]
-        cutoff = 2 if detector == 'threshold' else self.cutoff
+        if detector == 'threshold':
+            cutoff = 2
+        elif detector == 'click_counting':
+            cutoff = num_detectors + 1
+        elif detector == 'pnrd':
+            cutoff = self.cutoff
+        else:
+            raise ValueError(f'Unsupported detector: {detector}')
         if purity:
             for i in range(batch):
                 sample.append(_sample_pure(cov[i], mean[i], wires, self.nmode, cutoff, detector))
diff --git a/src/deepquantum/photonic/kensingtonian_.py b/src/deepquantum/photonic/kensingtonian_.py
index 0c682227..3eaa627e 100644
--- a/src/deepquantum/photonic/kensingtonian_.py
+++ b/src/deepquantum/photonic/kensingtonian_.py
@@ -1,82 +1,124 @@
-"""Functions for Kensingtonian"""
+"""Functions for Kensingtonian.
 
+The Kensingtonian appears in click-counting Gaussian boson sampling.
+See Phys. Rev. A 109, 023708 (2024), Eq. (27) and Appendix C.
+"""
+
+from collections import defaultdict
 from functools import cache
 from itertools import product
+from math import comb
 
 import torch
 
 
 @cache
-def _d_tuples(clicks: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
+def _d_term_data(
+    clicks: tuple[int, ...], num_detectors: int
+) -> tuple[tuple[tuple[int, ...], tuple[int, ...], tuple[float, ...], float], ...]:
+    size = 2 * len(clicks)
     ranges = [range(click + 1) for click in clicks]
-    return tuple(product(*ranges))
-
-
-def _multinomial_click_coeff(num_detectors: int, click: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
-    """Return N! / ((N-k)! (k-d)! d!) for each mode."""
-    n = torch.as_tensor(num_detectors, dtype=click.dtype, device=click.device)
-    return torch.exp(
-        torch.lgamma(n + 1) - torch.lgamma(n - click + 1) - torch.lgamma(click - d + 1) - torch.lgamma(d + 1)
-    )
-
+    data = []
+    for d_tuple in product(*ranges):
+        sign = -1 if sum(click - d for click, d in zip(clicks, d_tuple, strict=True)) % 2 else 1
+        coeff = 1
+        for click, d in zip(clicks, d_tuple, strict=True):
+            coeff *= comb(num_detectors, click) * comb(click, d)
+        keep_modes = tuple(i for i, d in enumerate(d_tuple) if d > 0)
+        prefactor = 1.0
+        for i in keep_modes:
+            prefactor *= num_detectors / d_tuple[i]
+        idx_tuple = tuple(sorted(keep_modes + tuple(i + size // 2 for i in keep_modes)))
+        dz_values = tuple((num_detectors - d_tuple[i]) / d_tuple[i] for i in keep_modes)
+        dz_values_tuple = dz_values + dz_values
+        data.append((d_tuple, idx_tuple, dz_values_tuple, sign * coeff * prefactor))
+    return tuple(data)
 
-def _mode_submat(mat: torch.Tensor, keep_modes: torch.Tensor) -> torch.Tensor:
-    """Keep both quadratures of each selected mode."""
-    if keep_modes.numel() == 0:
-        return mat.new_empty((0, 0))
-    idx1 = keep_modes
-    idx2 = idx1 + mat.shape[-1] // 2
-    idx = torch.sort(torch.cat([idx1, idx2]))[0]
-    return mat[idx][:, idx]
 
-
-def _mode_subvec(vec: torch.Tensor, keep_modes: torch.Tensor) -> torch.Tensor:
-    """Keep both quadratures of each selected mode from a vector."""
-    if keep_modes.numel() == 0:
-        return vec.new_empty((0,))
-    idx1 = keep_modes
-    idx2 = idx1 + vec.shape[-1] // 2
-    idx = torch.sort(torch.cat([idx1, idx2]))[0]
-    return vec[idx]
-
-
-def _dz_diag(d: torch.Tensor, keep_modes: torch.Tensor, num_detectors: int) -> torch.Tensor:
-    """Construct D_Z from Eq. (28) on the retained modes."""
-    values = (num_detectors - d[keep_modes]) / d[keep_modes]
-    return torch.cat([values, values]).diag()
+@cache
+def _grouped_term_data(
+    clicks: tuple[int, ...], num_detectors: int
+) -> tuple[
+    tuple[
+        int,
+        tuple[tuple[tuple[int, ...], tuple[float, ...], float], ...],
+    ],
+    ...,
+]:
+    """Group terms by submatrix size for batched linear algebra."""
+    groups = defaultdict(list)
+    for _, idx, dz_values, scalar in _d_term_data(clicks, num_detectors):
+        groups[len(idx)].append((idx, dz_values, scalar))
+    return tuple((size, tuple(groups[size])) for size in sorted(groups))
 
 
 def _kensingtonian_term(
     matrix: torch.Tensor,
     base_full: torch.Tensor,
-    clicks: torch.Tensor,
-    d: torch.Tensor,
-    num_detectors: int,
+    idx: torch.Tensor,
+    dz_values: torch.Tensor,
+    scalar: float,
     beta: torch.Tensor | None,
 ) -> torch.Tensor:
-    sign_power = torch.sum(clicks - d).to(torch.long)
-    sign = -1 if int(sign_power.item()) % 2 else 1
-
-    coeff = _multinomial_click_coeff(num_detectors, clicks, d).prod()
-    keep_modes = torch.nonzero(d > 0, as_tuple=False).flatten()
-    if keep_modes.numel() == 0:
+    if idx.numel() == 0:
         det_term = matrix.new_tensor(1)
         exp_term = matrix.new_tensor(1)
     else:
-        base = _mode_submat(base_full, keep_modes)
-        dz = _dz_diag(d, keep_modes, num_detectors).to(matrix.dtype)
-        det_mat = base + dz
+        base = base_full[idx][:, idx]
+        det_mat = base.clone()
+        det_mat.diagonal().add_(dz_values)
         det_term = torch.linalg.det(det_mat)
 
         if beta is None:
             exp_term = matrix.new_tensor(1)
         else:
-            rhs = _mode_subvec(beta, keep_modes)
+            rhs = beta[idx]
             quad = rhs @ torch.linalg.solve(det_mat, rhs)
             exp_term = torch.exp(quad)
 
-    prefactor = (num_detectors / d[keep_modes]).prod() if keep_modes.numel() > 0 else matrix.new_tensor(1)
-    return sign * coeff.to(matrix.dtype) * prefactor.to(matrix.dtype) * exp_term / torch.sqrt(det_term)
+    return matrix.new_tensor(scalar) * exp_term / torch.sqrt(det_term)
+
+
+def _kensingtonian_same_clicks(
+    matrix: torch.Tensor,
+    clicks: tuple[int, ...],
+    num_detectors: int,
+    alpha: torch.Tensor | None,
+) -> torch.Tensor:
+    """Evaluate a matrix batch whose samples share one click pattern."""
+    batch_size, size, _ = matrix.shape
+    identity = torch.eye(size, dtype=matrix.dtype, device=matrix.device)
+    base_full = identity - matrix
+    beta = None
+    if alpha is not None:
+        beta = (base_full @ alpha.unsqueeze(-1)).squeeze(-1)
+
+    result = matrix.new_zeros(batch_size)
+    for submatrix_size, terms in _grouped_term_data(clicks, num_detectors):
+        if submatrix_size == 0:
+            result = result + sum(term[2] for term in terms)
+            continue
+
+        # Creating these tensors once per group replaces per-term tensor
+        # construction and lets det/solve process all terms and samples at once.
+        idx = torch.tensor([term[0] for term in terms], dtype=torch.long, device=matrix.device)
+        dz_values = torch.tensor([term[1] for term in terms], dtype=matrix.dtype, device=matrix.device)
+        scalars = torch.tensor([term[2] for term in terms], dtype=matrix.dtype, device=matrix.device)
+
+        det_mat = base_full[:, idx[:, :, None], idx[:, None, :]]
+        det_mat.diagonal(dim1=-2, dim2=-1).add_(dz_values.unsqueeze(0))
+        det_term = torch.linalg.det(det_mat)
+
+        if beta is None:
+            numerator = scalars.unsqueeze(0)
+        else:
+            rhs = beta[:, idx]
+            solution = torch.linalg.solve(det_mat, rhs.unsqueeze(-1)).squeeze(-1)
+            quad = (rhs * solution).sum(dim=-1)
+            numerator = torch.exp(quad) * scalars.unsqueeze(0)
+
+        result = result + (numerator / torch.sqrt(det_term)).sum(dim=-1)
+    return result
 
 
 def kensingtonian(
@@ -85,7 +127,17 @@ def kensingtonian(
     num_detectors: int,
     alpha: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Calculate the Kensingtonian or loop Kensingtonian."""
+    """Calculate the Kensingtonian or loop Kensingtonian.
+
+    Args:
+        matrix: A 2M x 2M matrix A.
+        clicks: Length-M click pattern k, with 0 <= k_i <= num_detectors.
+        num_detectors: Number N of threshold detectors in each balanced click-counting detector.
+        alpha: Optional length-2M displacement vector for the loop Kensingtonian.
+
+    Returns:
+        Ken[A] for ``alpha is None`` and lken[A, alpha] otherwise.
+    """
     assert matrix.dim() == 2
     assert matrix.shape[-2] == matrix.shape[-1]
     assert matrix.shape[-1] % 2 == 0
@@ -97,20 +149,17 @@ def kensingtonian(
     assert torch.all(clicks >= 0)
     assert torch.all(clicks <= num_detectors)
 
-    clicks_float = clicks.to(matrix.real.dtype)
     if alpha is not None:
         assert alpha.shape == (2 * m,)
         alpha = alpha.to(dtype=matrix.dtype, device=matrix.device)
 
-    identity = torch.eye(matrix.shape[-1], dtype=matrix.dtype, device=matrix.device)
-    base_full = identity - matrix
-    beta = None if alpha is None else base_full @ alpha
-    ken = matrix.new_tensor(0)
-    clicks_tuple = tuple(int(click.item()) for click in clicks)
-    for d_tuple in _d_tuples(clicks_tuple):
-        d = torch.tensor(d_tuple, dtype=clicks_float.dtype, device=matrix.device)
-        ken = ken + _kensingtonian_term(matrix, base_full, clicks_float, d, num_detectors, beta)
-    return ken
+    clicks_tuple = tuple(clicks.tolist())
+    return _kensingtonian_same_clicks(
+        matrix.unsqueeze(0),
+        clicks_tuple,
+        num_detectors,
+        None if alpha is None else alpha.unsqueeze(0),
+    )[0]
 
 
 def kensingtonian_batch(
@@ -123,17 +172,36 @@ def kensingtonian_batch(
     assert matrix.dim() == 3
     assert matrix.shape[-2] == matrix.shape[-1]
     assert matrix.shape[-1] % 2 == 0
+    assert num_detectors >= 1
 
+    batch_size = matrix.shape[0]
+    m = matrix.shape[-1] // 2
     clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device)
     if clicks.dim() == 1:
-        clicks = clicks.expand(matrix.shape[0], -1)
+        clicks = clicks.expand(batch_size, -1)
+    assert clicks.shape == (batch_size, m)
+    assert torch.all(clicks >= 0)
+    assert torch.all(clicks <= num_detectors)
+
     if alpha is not None:
         alpha = torch.as_tensor(alpha, dtype=matrix.dtype, device=matrix.device)
         if alpha.dim() == 1:
-            alpha = alpha.expand(matrix.shape[0], -1)
+            alpha = alpha.expand(batch_size, -1)
+        assert alpha.shape == (batch_size, 2 * m)
+
+    # Group samples by click pattern. This keeps support for heterogeneous
+    # patterns while vectorizing the common homogeneous-batch case.
+    pattern_groups = defaultdict(list)
+    for sample, pattern in enumerate(clicks.tolist()):
+        pattern_groups[tuple(pattern)].append(sample)
 
+    indices = []
     values = []
-    for i in range(matrix.shape[0]):
-        alpha_i = None if alpha is None else alpha[i]
-        values.append(kensingtonian(matrix[i], clicks[i], num_detectors, alpha_i))
-    return torch.stack(values)
+    for pattern, sample_indices in pattern_groups.items():
+        index = torch.tensor(sample_indices, dtype=torch.long, device=matrix.device)
+        alpha_group = None if alpha is None else alpha.index_select(0, index)
+        values.append(_kensingtonian_same_clicks(matrix.index_select(0, index), pattern, num_detectors, alpha_group))
+        indices.append(index)
+
+    order = torch.cat(indices).argsort()
+    return torch.cat(values)[order]

From 20ae22eaa2e3b18c36ff7fa6fa84b9c2c35ab092 Mon Sep 17 00:00:00 2001
From: GetlinDu <1578440935@qq.com>
Date: Wed, 15 Jul 2026 15:38:27 +0800
Subject: [PATCH 4/4] Fix complex gamma handling in loop Kensingtonian

---
 src/deepquantum/photonic/circuit.py        |   4 +-
 src/deepquantum/photonic/kensingtonian_.py | 206 +++++++++++----------
 2 files changed, 106 insertions(+), 104 deletions(-)

diff --git a/src/deepquantum/photonic/circuit.py b/src/deepquantum/photonic/circuit.py
index 2acd5c74..16d03aa6 100644
--- a/src/deepquantum/photonic/circuit.py
+++ b/src/deepquantum/photonic/circuit.py
@@ -1204,8 +1204,8 @@ def _get_prob_gaussian_base(
             sub_mat = sub_matrix(matrix, final_state_double, final_state_double)
             prob = p_vac * torontonian(sub_mat, sub_gamma)
         elif detector == 'click_counting':
-            alpha = gamma if loop else None
-            prob = p_vac * kensingtonian(matrix, final_state, num_detectors, alpha)
+            loop_gamma = gamma if loop else None
+            prob = p_vac * kensingtonian(matrix, final_state, num_detectors, gamma=loop_gamma)
         return abs(prob.real).squeeze()
 
     def _get_prob_mps(self, final_state: Any, wires: int | list[int] | None = None) -> torch.Tensor:
diff --git a/src/deepquantum/photonic/kensingtonian_.py b/src/deepquantum/photonic/kensingtonian_.py
index 3eaa627e..a7654a70 100644
--- a/src/deepquantum/photonic/kensingtonian_.py
+++ b/src/deepquantum/photonic/kensingtonian_.py
@@ -1,8 +1,4 @@
-"""Functions for Kensingtonian.
-
-The Kensingtonian appears in click-counting Gaussian boson sampling.
-See Phys. Rev. A 109, 023708 (2024), Eq. (27) and Appendix C.
-"""
+"""Functions for real- and complex-representation Kensingtonians."""
 
 from collections import defaultdict
 from functools import cache
@@ -16,6 +12,7 @@
 def _d_term_data(
     clicks: tuple[int, ...], num_detectors: int
 ) -> tuple[tuple[tuple[int, ...], tuple[int, ...], tuple[float, ...], float], ...]:
+    """Return matrix-independent data for all mixed-radix ``d`` terms."""
     size = 2 * len(clicks)
     ranges = [range(click + 1) for click in clicks]
     data = []
@@ -30,8 +27,7 @@ def _d_term_data(
             prefactor *= num_detectors / d_tuple[i]
         idx_tuple = tuple(sorted(keep_modes + tuple(i + size // 2 for i in keep_modes)))
         dz_values = tuple((num_detectors - d_tuple[i]) / d_tuple[i] for i in keep_modes)
-        dz_values_tuple = dz_values + dz_values
-        data.append((d_tuple, idx_tuple, dz_values_tuple, sign * coeff * prefactor))
+        data.append((d_tuple, idx_tuple, dz_values + dz_values, sign * coeff * prefactor))
     return tuple(data)
 
 
@@ -52,46 +48,17 @@ def _grouped_term_data(
     return tuple((size, tuple(groups[size])) for size in sorted(groups))
 
 
-def _kensingtonian_term(
-    matrix: torch.Tensor,
-    base_full: torch.Tensor,
-    idx: torch.Tensor,
-    dz_values: torch.Tensor,
-    scalar: float,
-    beta: torch.Tensor | None,
-) -> torch.Tensor:
-    if idx.numel() == 0:
-        det_term = matrix.new_tensor(1)
-        exp_term = matrix.new_tensor(1)
-    else:
-        base = base_full[idx][:, idx]
-        det_mat = base.clone()
-        det_mat.diagonal().add_(dz_values)
-        det_term = torch.linalg.det(det_mat)
-
-        if beta is None:
-            exp_term = matrix.new_tensor(1)
-        else:
-            rhs = beta[idx]
-            quad = rhs @ torch.linalg.solve(det_mat, rhs)
-            exp_term = torch.exp(quad)
-
-    return matrix.new_tensor(scalar) * exp_term / torch.sqrt(det_term)
-
-
 def _kensingtonian_same_clicks(
     matrix: torch.Tensor,
     clicks: tuple[int, ...],
     num_detectors: int,
-    alpha: torch.Tensor | None,
+    gamma: torch.Tensor | None,
+    chunk_size: int | None,
 ) -> torch.Tensor:
     """Evaluate a matrix batch whose samples share one click pattern."""
     batch_size, size, _ = matrix.shape
     identity = torch.eye(size, dtype=matrix.dtype, device=matrix.device)
-    base_full = identity - matrix
-    beta = None
-    if alpha is not None:
-        beta = (base_full @ alpha.unsqueeze(-1)).squeeze(-1)
+    sigma_inv = identity - matrix
 
     result = matrix.new_zeros(batch_size)
     for submatrix_size, terms in _grouped_term_data(clicks, num_detectors):
@@ -99,98 +66,125 @@ def _kensingtonian_same_clicks(
             result = result + sum(term[2] for term in terms)
             continue
 
-        # Creating these tensors once per group replaces per-term tensor
-        # construction and lets det/solve process all terms and samples at once.
-        idx = torch.tensor([term[0] for term in terms], dtype=torch.long, device=matrix.device)
-        dz_values = torch.tensor([term[1] for term in terms], dtype=matrix.dtype, device=matrix.device)
-        scalars = torch.tensor([term[2] for term in terms], dtype=matrix.dtype, device=matrix.device)
-
-        det_mat = base_full[:, idx[:, :, None], idx[:, None, :]]
-        det_mat.diagonal(dim1=-2, dim2=-1).add_(dz_values.unsqueeze(0))
-        det_term = torch.linalg.det(det_mat)
-
-        if beta is None:
-            numerator = scalars.unsqueeze(0)
-        else:
-            rhs = beta[:, idx]
-            solution = torch.linalg.solve(det_mat, rhs.unsqueeze(-1)).squeeze(-1)
-            quad = (rhs * solution).sum(dim=-1)
-            numerator = torch.exp(quad) * scalars.unsqueeze(0)
-
-        result = result + (numerator / torch.sqrt(det_term)).sum(dim=-1)
+        group_chunk_size = len(terms) if chunk_size is None else chunk_size
+        for start in range(0, len(terms), group_chunk_size):
+            term_chunk = terms[start : start + group_chunk_size]
+            idx = torch.tensor(
+                [term[0] for term in term_chunk],
+                dtype=torch.long,
+                device=matrix.device,
+            )
+            dz_values = torch.tensor(
+                [term[1] for term in term_chunk],
+                dtype=matrix.dtype,
+                device=matrix.device,
+            )
+            scalars = torch.tensor(
+                [term[2] for term in term_chunk],
+                dtype=matrix.dtype,
+                device=matrix.device,
+            )
+
+            det_mat = sigma_inv[:, idx[:, :, None], idx[:, None, :]]
+            det_mat.diagonal(dim1=-2, dim2=-1).add_(dz_values.unsqueeze(0))
+            det_term = torch.linalg.det(det_mat)
+
+            if gamma is None:
+                numerator = scalars.unsqueeze(0)
+            else:
+                rhs = gamma[:, idx]
+                solution = torch.linalg.solve(det_mat, rhs.conj().unsqueeze(-1)).squeeze(-1)
+                quad = (rhs * solution).sum(dim=-1) / 2
+                numerator = torch.exp(quad) * scalars.unsqueeze(0)
+
+            result = result + (numerator / torch.sqrt(det_term)).sum(dim=-1)
     return result
 
 
 def kensingtonian(
     matrix: torch.Tensor,
-    clicks: torch.Tensor,
+    clicks: torch.Tensor | list[int],
     num_detectors: int,
-    alpha: torch.Tensor | None = None,
+    gamma: torch.Tensor | None = None,
+    chunk_size: int | None = None,
 ) -> torch.Tensor:
-    """Calculate the Kensingtonian or loop Kensingtonian.
+    r"""Calculate the Kensingtonian or loop Kensingtonian.
 
-    Args:
-        matrix: A 2M x 2M matrix A.
-        clicks: Length-M click pattern k, with 0 <= k_i <= num_detectors.
-        num_detectors: Number N of threshold detectors in each balanced click-counting detector.
-        alpha: Optional length-2M displacement vector for the loop Kensingtonian.
+    See https://arxiv.org/abs/2305.00853 Eq. (27)-(29) and Appendix C.
 
-    Returns:
-        Ken[A] for ``alpha is None`` and lken[A, alpha] otherwise.
+    Args:
+        matrix: The input matrix :math:`O=I-\Sigma^{-1}` in either
+            DeepQuantum's real ``xxpp`` ordering or complex ladder ordering.
+        clicks: The click-counting pattern :math:`k`.
+        num_detectors: The number of threshold detectors in each
+            click-counting detector.
+        gamma: The loop vector in the same convention as
+            :func:`torontonian`, namely
+            :math:`\gamma=\sqrt{2}\alpha^T\Sigma^{-1}`.
+            Default: ``None``
+        chunk_size: The maximum number of ``d`` terms evaluated together
+            within one same-size group. Default: ``None``
     """
-    assert matrix.dim() == 2
-    assert matrix.shape[-2] == matrix.shape[-1]
-    assert matrix.shape[-1] % 2 == 0
-    assert num_detectors >= 1
+    assert matrix.dim() == 2, 'Input matrix should be 2D'
+    assert matrix.shape[-2] == matrix.shape[-1], 'Input matrix should be square'
+    assert matrix.shape[-1] % 2 == 0, 'Input matrix dimension should be even'
+    assert num_detectors >= 1, 'num_detectors should be positive'
+    assert chunk_size is None or chunk_size > 0, 'chunk_size should be positive'
 
-    m = matrix.shape[-1] // 2
+    nmode = matrix.shape[-1] // 2
     clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device)
-    assert clicks.shape == (m,)
-    assert torch.all(clicks >= 0)
-    assert torch.all(clicks <= num_detectors)
+    clicks = clicks.reshape(-1)
+    assert clicks.shape == (nmode,), f'Click pattern should have shape ({nmode},)'
+    assert (clicks >= 0).all(), 'Click numbers should be non-negative'
+    assert (clicks <= num_detectors).all(), 'Click numbers should not exceed num_detectors'
 
-    if alpha is not None:
-        assert alpha.shape == (2 * m,)
-        alpha = alpha.to(dtype=matrix.dtype, device=matrix.device)
+    if gamma is not None:
+        gamma = torch.as_tensor(gamma, device=matrix.device)
+        assert matrix.is_complex() or not gamma.is_complex(), 'Complex gamma requires a complex matrix'
+        gamma = gamma.to(matrix.dtype).reshape(-1)
+        assert gamma.shape == (2 * nmode,), f'gamma should have shape ({2 * nmode},)'
 
     clicks_tuple = tuple(clicks.tolist())
     return _kensingtonian_same_clicks(
         matrix.unsqueeze(0),
         clicks_tuple,
         num_detectors,
-        None if alpha is None else alpha.unsqueeze(0),
+        None if gamma is None else gamma.unsqueeze(0),
+        chunk_size,
     )[0]
 
 
 def kensingtonian_batch(
     matrix: torch.Tensor,
-    clicks: torch.Tensor,
+    clicks: torch.Tensor | list[int],
     num_detectors: int,
-    alpha: torch.Tensor | None = None,
+    gamma: torch.Tensor | None = None,
+    chunk_size: int | None = None,
 ) -> torch.Tensor:
-    """Calculate batched Kensingtonians."""
-    assert matrix.dim() == 3
-    assert matrix.shape[-2] == matrix.shape[-1]
-    assert matrix.shape[-1] % 2 == 0
-    assert num_detectors >= 1
+    """Calculate batched Kensingtonians for one or more click patterns."""
+    assert matrix.dim() == 3, 'Input tensor should be in batched size'
+    assert matrix.shape[-2] == matrix.shape[-1], 'Input matrices should be square'
+    assert matrix.shape[-1] % 2 == 0, 'Input matrix dimension should be even'
+    assert num_detectors >= 1, 'num_detectors should be positive'
+    assert chunk_size is None or chunk_size > 0, 'chunk_size should be positive'
 
     batch_size = matrix.shape[0]
-    m = matrix.shape[-1] // 2
+    nmode = matrix.shape[-1] // 2
     clicks = torch.as_tensor(clicks, dtype=torch.long, device=matrix.device)
     if clicks.dim() == 1:
         clicks = clicks.expand(batch_size, -1)
-    assert clicks.shape == (batch_size, m)
-    assert torch.all(clicks >= 0)
-    assert torch.all(clicks <= num_detectors)
-
-    if alpha is not None:
-        alpha = torch.as_tensor(alpha, dtype=matrix.dtype, device=matrix.device)
-        if alpha.dim() == 1:
-            alpha = alpha.expand(batch_size, -1)
-        assert alpha.shape == (batch_size, 2 * m)
-
-    # Group samples by click pattern. This keeps support for heterogeneous
-    # patterns while vectorizing the common homogeneous-batch case.
+    assert clicks.shape == (batch_size, nmode), f'Click patterns should have shape ({batch_size}, {nmode})'
+    assert (clicks >= 0).all(), 'Click numbers should be non-negative'
+    assert (clicks <= num_detectors).all(), 'Click numbers should not exceed num_detectors'
+
+    if gamma is not None:
+        gamma = torch.as_tensor(gamma, device=matrix.device)
+        assert matrix.is_complex() or not gamma.is_complex(), 'Complex gamma requires complex matrices'
+        gamma = gamma.to(matrix.dtype)
+        if gamma.dim() == 1:
+            gamma = gamma.expand(batch_size, -1)
+        assert gamma.shape == (batch_size, 2 * nmode), f'gamma should have shape ({batch_size}, {2 * nmode})'
+
     pattern_groups = defaultdict(list)
     for sample, pattern in enumerate(clicks.tolist()):
         pattern_groups[tuple(pattern)].append(sample)
@@ -199,8 +193,16 @@ def kensingtonian_batch(
     values = []
     for pattern, sample_indices in pattern_groups.items():
         index = torch.tensor(sample_indices, dtype=torch.long, device=matrix.device)
-        alpha_group = None if alpha is None else alpha.index_select(0, index)
-        values.append(_kensingtonian_same_clicks(matrix.index_select(0, index), pattern, num_detectors, alpha_group))
+        gamma_group = None if gamma is None else gamma.index_select(0, index)
+        values.append(
+            _kensingtonian_same_clicks(
+                matrix.index_select(0, index),
+                pattern,
+                num_detectors,
+                gamma_group,
+                chunk_size,
+            )
+        )
         indices.append(index)
 
     order = torch.cat(indices).argsort()