NaN gradients in Gaussian backend probability backward for mixed init_state

[Jooyuza]

Description

A finite forward probability produces NaN gradients during backward() in the Gaussian backend.

The forward probability is finite, but after calling backward(), all trainable parameters of two BeamSplitter gates have NaN gradients.

Minimal Reproduction

import torch
import deepquantum as dq

cov = torch.eye(8)
cov[1, 1] = 2.
cov[5, 5] = 2.
thermal = dq.GaussianState([cov, torch.zeros(1, 8, 1)], cutoff=3)

cir = dq.QumodeCircuit(nmode=4, init_state=thermal, backend="gaussian", cutoff=3)
cir.bs([0, 1])
cir.bs([2, 3])

prob = next(iter(cir(is_prob=True, sort=False).values())).sum()
print("forward:", prob)

prob.backward()
for name, param in cir.named_parameters():
    print(name, param.grad)

Actual Output

forward: tensor(0.6667, grad_fn=<SumBackward0>)
operators.0.theta tensor(nan)
operators.0.phi tensor(nan)
operators.1.theta tensor(nan)
operators.1.phi tensor(nan)

Expected Behavior

The forward probability is finite, so the gradients of the BeamSplitter parameters should also be finite, or the operation should raise a clear numerical/domain error before returning NaN gradients.

Environment

torch 2.4.0
cuda available True
deepquantum 4.5.0
backend="gaussian"
detector="pnrd" default

The reproduction uses CPU tensors even though CUDA is available.

Extra Debug Info

With torch.autograd.set_detect_anomaly(True), PyTorch reports:

RuntimeError: Function 'ProdBackward1' returned nan values in its 0th output.

The traceback points to the hafnian calculation:

deepquantum/photonic/hafnian_.py:88
poly_prod = poly_list.prod()

called from:

deepquantum/photonic/circuit.py:1171
haf = abs(hafnian(sub_mat, loop=loop)) ** 2 if purity else hafnian(sub_mat, loop=loop)

[Jooyuza]

It's interesting that torch 2.3.1 works fine

[Jooyuza]

问题出在 hafnian_.py (line 39) 里 poly_list.prod()：在 Gaussian 概率的 vmap + hafnian 反传图中，零乘积点会让 ProdBackward 返回 nan。
如果用一个简单stable_prod替换可以跑通

poly_prod = stable_prod(poly_list)

def stable_prod(tensor: torch.Tensor) -> torch.Tensor:
    """Multiply entries without using ``Tensor.prod`` to avoid NaN gradients at zeros."""
    result = tensor.new_ones(())
    for value in tensor:
        result = result * value
    return result