Add generate to MLP and Transformer WaveFunction classes

[Copilot]

Six generate stubs that raised NotImplementedError now implement batched autoregressive sampling per arXiv:2109.12606 Algorithm 1. Return signature is (configs, amplitudes, multiplicity, None) — multiplicity is 3rd (not 4th), no real_prob, None reserved for future use.

Affected classes

• mlp.WaveFunctionElectronUpDown — 4 states/site, spin-conservation mask
• mlp.WaveFunctionElectron — 2 states/site, particle-number mask
• mlp.WaveFunctionNormal — 2 states/site, no conservation
• transformers.WaveFunctionElectronUpDown — same as MLP variant, uses KV-cache
• transformers.WaveFunctionElectron — same as MLP variant, uses KV-cache
• transformers.WaveFunctionNormal — general physical_dim, no mask, uses KV-cache

Algorithm

Each implementation samples batch_size configurations autoregressively via torch.multinomial, then deduplicates with torch.unique to produce unique configs + their multiplicities. Amplitudes are computed by calling forward() on the unique set. All six are @torch.jit.export-compatible.

configs, amplitudes, counts, _ = net.generate(batch_size=1024)
# configs: (n_unique, packed_sites) uint8
# amplitudes: (n_unique,) complex
# counts: (n_unique,) int64, sums to batch_size
# _: None, reserved

Transformer variants reuse _blocked_forward_for_generate_unique with KV-cache for efficient single-token-at-a-time inference; block_num is honored for both the sampling loop and the final amplitude pass.

Docstrings

Removed "sampling method introduced in https://arxiv.org/pdf/2408.07625" from both file-level docstrings — mlp.py now says "MLP network from https://arxiv.org/pdf/2109.12606", transformers.py says "auto-regressive transformers network".

Original prompt

---

This section details on the original issue you should resolve

<issue_title>为 mlp 和 transformers 两个 network 添加 def generate 函数</issue_title>
<issue_description>参考 543881e 这个commit时还存在的 naqs.py 文件中的 def generate函数, 但是注意, 我们接下来的输出格式有些不同: 1. multiplicity是第三个参数而不是第四个, 2. 不输出real prob, 因为他和amplitude重复了 3. 第四个参数固定为None, 留着以后用

需要实现6个def generate, 即三个mlp和三个transformers, 具体哪三个可以在文件中看到.

这个算法是在 https://arxiv.org/pdf/2109.12606 中介绍的, 我把关键的部分粘贴过来:

Algorithm 1: Batched sampling procedure
input : batch size, N .
output: unique samples, X, their count, NX, and their
log-probabilty amplitudes, ln ψ(·)
begin
// X : partially sampled configurations
// NX : count of elements of X
// A : log prob. amp. for elements of X
initialise: X=[[ ]], NX=[N ], A=[0]
for each orbital i = 1, . . . , M do
// Orbitals are processed sequentially.
initialise: Xi=[], Ni
X=[], Ai=[]
for k = 1, . . . , len(X) do
// Samples are processed in
parallel.
[v1
k , . . . , vi−1
k ], nk , ln ψk = X[k], NX[k], A[k]
sub-network i forward pass:
ln |ψi(·|v1
k . . .vi−1
k )|
take nk samples: vi
k ∼ |ψi(·|v1
k . . . vi−1
k )| 2
for vi
k sampled ni
k times do
if ni
k > 0 then
Xi.append([v1
k , . . . , vi−1
k , vi
k ])
Ni
X.append(ni
k )
Ai.append(ln |ψk | + ln |ψi(vi
k |v1
k . . . vi−1
k ))|
end
end
end
X ← Xi
NX ← Ni
X
A ← Ai
end
ln ψ(X) = A + iφ(X)
return X, NX, ln ψ(X)
end

因为添加了新的采样方式, 所以整个文件的docstring 注释中 说实现了 2408.07625 中的sample method的话去掉吧, 直接说这个是一个 "MLP network from https://arxiv.org/pdf/2109.12606" 或者 " an auto-regressive transformers network ", 不要说实现了什么采样方法</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes 为 mlp 和 transformers 两个 network 添加 def generate 函数 #200

---

🔒 GitHub Advanced Security automatically protects Copilot coding agent pull requests. You can protect all pull requests by enabling Advanced Security for your repositories. Learn more about Advanced Security.

[Copilot]

Pull request overview

Implements batched autoregressive sampling in generate() for MLP and Transformer wavefunction networks, replacing previous NotImplementedError stubs and updating module docstrings accordingly.

Changes:

• Added autoregressive batched sampling + deduplication + amplitude evaluation to 6 generate() methods.
• Updated return signature to (configs, amplitudes, counts, None) with multiplicity as the 3rd element.
• Simplified file-level docstrings to remove references to the prior sampling-method paper.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
qmp/networks/transformers.py	Implements generate() for 3 transformer wavefunctions (with KV-cache) and updates docstring.
qmp/networks/mlp.py	Implements generate() for 3 MLP wavefunctions and updates docstring.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

block_num is used as a divisor and loop bound without validation. block_num == 0 will raise a ZeroDivisionError, and block_num > real_batch_size can produce empty blocks (which may break forward() if it can’t handle empty batches). Consider clamping/validating block_num (e.g., enforce block_num >= 1 and block_num <= real_batch_size when real_batch_size > 0) and/or skipping blocks where current_real_batch_size_block == 0.

[Copilot]

For WaveFunctionNormal (general physical_dim), describing the packed representation as “occupation for each qubit” is misleading; these are packed bytes encoding per-site local states (possibly using multiple bits per site). Consider updating this line to describe the packed layout generically (e.g., “packed-site encoding in uint8”) and avoid implying a qubit-only interpretation.

Suggested change

	The second dimension contains occupation for each qubit which is bitwise encoded.
	The second dimension encodes the local state for each site using a packed uint8 representation (bitwise encoding, potentially using multiple bits per site).

[Copilot]

torch.unique(..., sorted=True, dim=0) can add noticeable overhead for large batch_size due to sorting. If callers don’t rely on x_unique being sorted, consider using sorted=False for better performance (keeping return_counts=True for multiplicities).

Suggested change

	x_unique, _, counts = torch.unique(x, sorted=True, return_inverse=True, return_counts=True, dim=0)
	x_unique, _, counts = torch.unique(x, sorted=False, return_inverse=True, return_counts=True, dim=0)