From 99632fc0e5215ddaad3891f989b59bb133c6074f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 5 Mar 2026 19:11:49 +0000 Subject: [PATCH 1/4] Initial plan From 8225a6d56c23cae73460cea7d710423cbbf73c82 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 5 Mar 2026 19:26:32 +0000 Subject: [PATCH 2/4] Add MPS network and generate() method to NetworkProto and existing networks Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com> --- qmp/networks/mlp.py | 21 ++++ qmp/networks/mps.py | 217 +++++++++++++++++++++++++++++++++++ qmp/networks/transformers.py | 21 ++++ qmp/utility/model_dict.py | 22 ++++ 4 files changed, 281 insertions(+) create mode 100644 qmp/networks/mps.py diff --git a/qmp/networks/mlp.py b/qmp/networks/mlp.py index 50ef000..265ab05 100644 --- a/qmp/networks/mlp.py +++ b/qmp/networks/mlp.py @@ -330,6 +330,13 @@ def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Te return x, amplitude, None, None + @torch.jit.export + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate is not implemented for this network. + """ + raise RuntimeError("generate is not implemented for this network") + class WaveFunctionElectron(torch.nn.Module): """ @@ -568,6 +575,13 @@ def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Te return x, amplitude, None, None + @torch.jit.export + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate is not implemented for this network. + """ + raise RuntimeError("generate is not implemented for this network") + class WaveFunctionNormal(torch.nn.Module): """ @@ -778,3 +792,10 @@ def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Te amplitude: torch.Tensor = torch.cat(amplitude_list) return x, amplitude, None, None + + @torch.jit.export + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate is not implemented for this network. + """ + raise RuntimeError("generate is not implemented for this network") diff --git a/qmp/networks/mps.py b/qmp/networks/mps.py new file mode 100644 index 0000000..43b8dfa --- /dev/null +++ b/qmp/networks/mps.py @@ -0,0 +1,217 @@ +""" +This file implements a trivial Matrix Product State (MPS) network. +""" + +import torch +from ..utility.bitspack import pack_int, unpack_int + + +class WaveFunction(torch.nn.Module): + """ + A trivial Matrix Product State (MPS) wave function. + + The MPS tensors have shape (physical_dim, virtual_dim, virtual_dim) for each site, + with open boundary conditions: the left boundary vector is e_0 = [1, 0, ..., 0] + and the right boundary vector is e_0 = [1, 0, ..., 0]. + + The amplitude is computed as: + Ψ(s) = e_L^T @ A[0][s_0] @ A[1][s_1] @ ... @ A[L-1][s_{L-1}] @ e_R + where e_L = e_R = [1, 0, ..., 0] (first standard basis vector). + """ + + def __init__( + self, + *, + sites: int, # Number of sites in the MPS + physical_dim: int, # Dimension of the physical (local) Hilbert space at each site + virtual_dim: int, # Bond dimension (virtual dimension) of the MPS + ) -> None: + super().__init__() + self.sites: int = sites + self.physical_dim: int = physical_dim + self.virtual_dim: int = virtual_dim + + # MPS tensors: (sites, physical_dim, virtual_dim, virtual_dim) + # All sites use the same tensor shape; open boundary conditions are imposed + # by contracting with the first standard basis vector e_0 on both ends. + self.tensors: torch.nn.Parameter = torch.nn.Parameter( + torch.randn(sites, physical_dim, virtual_dim, virtual_dim) + ) + + # Dummy Parameter for Device and Dtype Retrieval + # This parameter is used to infer the device and dtype of the model. + self.dummy_param = torch.nn.Parameter(torch.empty(0)) + + @torch.jit.export + def _bit_size(self) -> int: + """ + Return the number of bits needed to encode a single physical state. + """ + if self.physical_dim <= 1 << 1: + return 1 + if self.physical_dim <= 1 << 2: + return 2 + if self.physical_dim <= 1 << 4: + return 4 + if self.physical_dim <= 1 << 8: + return 8 + raise ValueError("physical_dim should be less than or equal to 256") + + @torch.jit.export + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Compute the wave function amplitude for the given configurations. + + Parameters + ---------- + x : torch.Tensor + Packed configurations as a 2D uint8 tensor of shape (batch_size, packed_sites). + + Returns + ------- + torch.Tensor + Complex128 amplitudes of shape (batch_size,). + """ + device: torch.device = self.dummy_param.device + dtype: torch.dtype = self.dummy_param.dtype + + batch_size: int = x.shape[0] + # x_unpacked: batch_size * sites, each value in [0, physical_dim) + x_unpacked: torch.Tensor = unpack_int(x, size=self._bit_size(), last_dim=self.sites).view( + [batch_size, self.sites] + ) + + tensors: torch.Tensor = self.tensors.to(dtype=dtype) + + # v: batch_size * virtual_dim -- left environment vector, starting from e_L = [1, 0, ..., 0] + v: torch.Tensor = torch.zeros([batch_size, self.virtual_dim], device=device, dtype=dtype) + v[:, 0] = 1.0 + + for i in range(self.sites): + # A_i: physical_dim * virtual_dim * virtual_dim + A_i: torch.Tensor = tensors[i] + # s_i: batch_size, integer indices for the physical state at site i + s_i: torch.Tensor = x_unpacked[:, i].to(dtype=torch.int64) + # A_s_i: batch_size * virtual_dim * virtual_dim + A_s_i: torch.Tensor = A_i[s_i] + # v = v @ A_s_i: (batch_size, virtual_dim) + v = torch.bmm(v.unsqueeze(1), A_s_i).squeeze(1) + + # Project onto right boundary e_R = [1, 0, ..., 0]: extract the first element + amplitude: torch.Tensor = v[:, 0] + + # Return as complex128 tensor with zero imaginary part + return torch.view_as_complex( + torch.stack([amplitude.double(), torch.zeros_like(amplitude).double()], dim=-1) + ) + + @torch.jit.export + def generate_unique( + self, batch_size: int, block_num: int = 1 + ) -> tuple[torch.Tensor, torch.Tensor, None, None]: + """ + Unique generation is not implemented for the MPS network. + """ + raise RuntimeError("generate_unique is not implemented for the MPS network") + + @torch.jit.export + def generate( + self, batch_size: int, block_num: int = 1 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate configurations by sampling from the MPS distribution. + + Configurations are sampled autoregressively from left to right using + right-environment density matrices. After sampling, identical configurations + are grouped and counted. + + Parameters + ---------- + batch_size : int + The number of configurations to sample (before deduplication). + block_num : int, default=1 + Currently unused; kept for interface consistency. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor, None] + - configs : unique packed configurations, shape (n_unique, packed_sites), uint8. + - amplitudes : wave function amplitudes for the unique configs, shape (n_unique,), complex128. + - counts : number of times each unique config was sampled, shape (n_unique,), int64. + - None + """ + device: torch.device = self.dummy_param.device + dtype: torch.dtype = self.dummy_param.dtype + + tensors: torch.Tensor = self.tensors.to(dtype=dtype) + + # ------------------------------------------------------------------ # + # Pre-compute right-environment density matrices from right to left. # + # # + # renv[k] = M[sites - k], where: # + # M[sites] = e_R @ e_R^T (right-boundary density matrix) # + # M[i] = sum_s A[i][s] @ M[i+1] @ A[i][s]^T for i < sites # + # # + # When sampling at site i we need M[i+1] = renv[sites - i - 1]. # + # ------------------------------------------------------------------ # + renv: list[torch.Tensor] = [] + M: torch.Tensor = torch.zeros([self.virtual_dim, self.virtual_dim], device=device, dtype=dtype) + M[0, 0] = 1.0 # e_R @ e_R^T (only the [0,0] entry is non-zero) + renv.append(M) + + for i in range(self.sites - 1, 0, -1): + A_i: torch.Tensor = tensors[i] # physical_dim * virtual_dim * virtual_dim + # M_new[a,d] = sum_{s,b,c} A_i[s,a,b] * M[b,c] * A_i[s,d,c] + # = sum_s (A_i[s] @ M @ A_i[s]^T)[a,d] + M = torch.einsum("sab,bc,sdc->ad", A_i, M, A_i) + renv.append(M) + + # After the loop renv has length `sites`: + # renv[0] = M[sites], renv[1] = M[sites-1], ..., renv[sites-1] = M[1] + + # ------------------------------------------------------------------ # + # Autoregressive sampling. # + # ------------------------------------------------------------------ # + # v: batch_size * virtual_dim (left state, initialised to e_L) + v: torch.Tensor = torch.zeros([batch_size, self.virtual_dim], device=device, dtype=dtype) + v[:, 0] = 1.0 + + samples: torch.Tensor = torch.zeros([batch_size, self.sites], device=device, dtype=torch.uint8) + + for i in range(self.sites): + A_i = tensors[i] # physical_dim * virtual_dim * virtual_dim + # M_next is the right-environment density matrix for the positions after site i + M_next: torch.Tensor = renv[self.sites - i - 1] + + # v_new_all[b, s, j] = sum_k v[b, k] * A_i[s, k, j] (batch, phys, virt) + v_new_all: torch.Tensor = torch.einsum("bi,sij->bsj", v, A_i) + + # prob[b, s] = v_new_all[b, s] @ M_next @ v_new_all[b, s]^T (non-negative) + Mv: torch.Tensor = torch.einsum("jk,bsk->bsj", M_next, v_new_all) + prob: torch.Tensor = (v_new_all * Mv).sum(dim=-1) # batch_size * physical_dim + + # Correct any numerical round-off that may produce tiny negative values + prob = prob.clamp(min=0.0) + + # Sample a physical state at site i for each batch element + s_i: torch.Tensor = torch.multinomial(prob, num_samples=1).squeeze(1) # batch_size + samples[:, i] = s_i.to(torch.uint8) + + # Advance the left state vector + arange: torch.Tensor = torch.arange(batch_size, device=device) + v = v_new_all[arange, s_i, :] # batch_size * virtual_dim + + # ------------------------------------------------------------------ # + # Pack, deduplicate, and count. # + # ------------------------------------------------------------------ # + x_packed: torch.Tensor = pack_int(samples, size=self._bit_size()) + + # torch.unique returns (unique, inverse_indices, counts) in TorchScript + x_unique, _, counts = torch.unique( + x_packed, sorted=True, return_inverse=True, return_counts=True, dim=0 + ) + + # Compute amplitudes for the unique configurations + amplitudes: torch.Tensor = self(x_unique) + + return x_unique, amplitudes, counts, None diff --git a/qmp/networks/transformers.py b/qmp/networks/transformers.py index e5fbd48..08f5ce4 100644 --- a/qmp/networks/transformers.py +++ b/qmp/networks/transformers.py @@ -626,6 +626,13 @@ def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Te return x, amplitude, None, None + @torch.jit.export + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate is not implemented for this network. + """ + raise RuntimeError("generate is not implemented for this network") + class WaveFunctionElectron(torch.nn.Module): """ @@ -969,6 +976,13 @@ def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Te return x, amplitude, None, None + @torch.jit.export + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate is not implemented for this network. + """ + raise RuntimeError("generate is not implemented for this network") + class WaveFunctionNormal(torch.nn.Module): """ @@ -1280,3 +1294,10 @@ def _bit_size(self) -> int: if self.physical_dim <= 1 << 8: return 8 raise ValueError("physical_dim should be less than or equal to 256") + + @torch.jit.export + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate is not implemented for this network. + """ + raise RuntimeError("generate is not implemented for this network") diff --git a/qmp/utility/model_dict.py b/qmp/utility/model_dict.py index 883aa4d..4d9aa6a 100644 --- a/qmp/utility/model_dict.py +++ b/qmp/utility/model_dict.py @@ -52,6 +52,28 @@ def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Te The last two None values are reserved for future use. """ + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + """ + Generate a batch of configurations (possibly with repetitions). + + Parameters + ---------- + batch_size : int + The number of configurations to sample before deduplication. + block_num : int, default=1 + The number of batch blocks to use. It is used to split the batch into smaller parts to avoid memory issues. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor, None] + A tuple containing the unique configurations, their amplitudes, their sample counts, and a None value. + The configurations are a two-dimensional uint8 tensor. + The second dimension contains occupation for each qubit which is bitwise encoded. + The amplitudes are a one-dimensional complex tensor with the only dimension equal to the number of unique configurations. + The counts are a one-dimensional int64 tensor recording how many times each unique configuration was sampled. + The last None value is reserved for future use. + """ + def load_state_dict(self, data: dict[str, torch.Tensor]) -> typing.Any: """torch.nn.Module function""" From 03984c123b1d014db4b1ebd62861d8fc05447b35 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 5 Mar 2026 19:47:51 +0000 Subject: [PATCH 3/4] Use NotImplementedError and align generate docstring style with generate_unique Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com> --- qmp/networks/mlp.py | 6 +++--- qmp/networks/mps.py | 24 +++++++++++------------- qmp/networks/transformers.py | 6 +++--- qmp/utility/model_dict.py | 10 +++++----- 4 files changed, 22 insertions(+), 24 deletions(-) diff --git a/qmp/networks/mlp.py b/qmp/networks/mlp.py index 265ab05..494bd55 100644 --- a/qmp/networks/mlp.py +++ b/qmp/networks/mlp.py @@ -335,7 +335,7 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t """ Generate is not implemented for this network. """ - raise RuntimeError("generate is not implemented for this network") + raise NotImplementedError("generate is not implemented for this network") class WaveFunctionElectron(torch.nn.Module): @@ -580,7 +580,7 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t """ Generate is not implemented for this network. """ - raise RuntimeError("generate is not implemented for this network") + raise NotImplementedError("generate is not implemented for this network") class WaveFunctionNormal(torch.nn.Module): @@ -798,4 +798,4 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t """ Generate is not implemented for this network. """ - raise RuntimeError("generate is not implemented for this network") + raise NotImplementedError("generate is not implemented for this network") diff --git a/qmp/networks/mps.py b/qmp/networks/mps.py index 43b8dfa..b5f4319 100644 --- a/qmp/networks/mps.py +++ b/qmp/networks/mps.py @@ -110,35 +110,33 @@ def generate_unique( self, batch_size: int, block_num: int = 1 ) -> tuple[torch.Tensor, torch.Tensor, None, None]: """ - Unique generation is not implemented for the MPS network. + Generate a batch of unique configurations. """ - raise RuntimeError("generate_unique is not implemented for the MPS network") + raise NotImplementedError("generate_unique is not implemented for the MPS network") @torch.jit.export def generate( self, batch_size: int, block_num: int = 1 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: """ - Generate configurations by sampling from the MPS distribution. - - Configurations are sampled autoregressively from left to right using - right-environment density matrices. After sampling, identical configurations - are grouped and counted. + Generate a batch of configurations. Parameters ---------- batch_size : int - The number of configurations to sample (before deduplication). + The number of configurations to generate. block_num : int, default=1 - Currently unused; kept for interface consistency. + The number of batch block to generate. It is used to split the batch into smaller parts to avoid memory issues. Returns ------- tuple[torch.Tensor, torch.Tensor, torch.Tensor, None] - - configs : unique packed configurations, shape (n_unique, packed_sites), uint8. - - amplitudes : wave function amplitudes for the unique configs, shape (n_unique,), complex128. - - counts : number of times each unique config was sampled, shape (n_unique,), int64. - - None + A tuple containing the generated configurations, their amplitudes, their sample counts, and a None value. + The configurations are a two-dimensional uint8 tensor with first dimension equal to the number of unique configurations. + The second dimension contains occupation for each qubit which is bitwise encoded. + The amplitudes are a one-dimensional complex tensor with the only dimension equal to the number of unique configurations. + The counts are a one-dimensional int64 tensor recording how many times each unique configuration was sampled. + The last None value is reserved for future use. """ device: torch.device = self.dummy_param.device dtype: torch.dtype = self.dummy_param.dtype diff --git a/qmp/networks/transformers.py b/qmp/networks/transformers.py index 08f5ce4..507a11a 100644 --- a/qmp/networks/transformers.py +++ b/qmp/networks/transformers.py @@ -631,7 +631,7 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t """ Generate is not implemented for this network. """ - raise RuntimeError("generate is not implemented for this network") + raise NotImplementedError("generate is not implemented for this network") class WaveFunctionElectron(torch.nn.Module): @@ -981,7 +981,7 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t """ Generate is not implemented for this network. """ - raise RuntimeError("generate is not implemented for this network") + raise NotImplementedError("generate is not implemented for this network") class WaveFunctionNormal(torch.nn.Module): @@ -1300,4 +1300,4 @@ def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, t """ Generate is not implemented for this network. """ - raise RuntimeError("generate is not implemented for this network") + raise NotImplementedError("generate is not implemented for this network") diff --git a/qmp/utility/model_dict.py b/qmp/utility/model_dict.py index 4d9aa6a..2df8bb7 100644 --- a/qmp/utility/model_dict.py +++ b/qmp/utility/model_dict.py @@ -54,20 +54,20 @@ def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Te def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: """ - Generate a batch of configurations (possibly with repetitions). + Generate a batch of configurations. Parameters ---------- batch_size : int - The number of configurations to sample before deduplication. + The number of configurations to generate. block_num : int, default=1 - The number of batch blocks to use. It is used to split the batch into smaller parts to avoid memory issues. + The number of batch block to generate. It is used to split the batch into smaller parts to avoid memory issues. Returns ------- tuple[torch.Tensor, torch.Tensor, torch.Tensor, None] - A tuple containing the unique configurations, their amplitudes, their sample counts, and a None value. - The configurations are a two-dimensional uint8 tensor. + A tuple containing the generated configurations, their amplitudes, their sample counts, and a None value. + The configurations are a two-dimensional uint8 tensor with first dimension equal to the number of unique configurations. The second dimension contains occupation for each qubit which is bitwise encoded. The amplitudes are a one-dimensional complex tensor with the only dimension equal to the number of unique configurations. The counts are a one-dimensional int64 tensor recording how many times each unique configuration was sampled. From c8584d749c4d8885f4947d4c7a786d973295e428 Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Fri, 6 Mar 2026 03:58:55 +0800 Subject: [PATCH 4/4] Ruff format. --- qmp/networks/mps.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/qmp/networks/mps.py b/qmp/networks/mps.py index b5f4319..6eecc0f 100644 --- a/qmp/networks/mps.py +++ b/qmp/networks/mps.py @@ -101,23 +101,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: amplitude: torch.Tensor = v[:, 0] # Return as complex128 tensor with zero imaginary part - return torch.view_as_complex( - torch.stack([amplitude.double(), torch.zeros_like(amplitude).double()], dim=-1) - ) + return torch.view_as_complex(torch.stack([amplitude.double(), torch.zeros_like(amplitude).double()], dim=-1)) @torch.jit.export - def generate_unique( - self, batch_size: int, block_num: int = 1 - ) -> tuple[torch.Tensor, torch.Tensor, None, None]: + def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, None, None]: """ Generate a batch of unique configurations. """ raise NotImplementedError("generate_unique is not implemented for the MPS network") @torch.jit.export - def generate( - self, batch_size: int, block_num: int = 1 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: + def generate(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, None]: """ Generate a batch of configurations. @@ -205,9 +199,7 @@ def generate( x_packed: torch.Tensor = pack_int(samples, size=self._bit_size()) # torch.unique returns (unique, inverse_indices, counts) in TorchScript - x_unique, _, counts = torch.unique( - x_packed, sorted=True, return_inverse=True, return_counts=True, dim=0 - ) + x_unique, _, counts = torch.unique(x_packed, sorted=True, return_inverse=True, return_counts=True, dim=0) # Compute amplitudes for the unique configurations amplitudes: torch.Tensor = self(x_unique)