🐛 contract in some orders on a circuit with ring struction may get incorrect multiplication counts.

[Yonv1943]

收缩 sycamore 以及 tensor grid ，tensor ring 这种有环状结构的电路，会有bug，导致乘法次数计算错误
（刚好我们测试的 tensor train，tensor tree 没有环状结构）

需要有环状结构，且按某个顺序收缩张量节点，才会触发

下面的代码，在一个小规模的 sycamore 电路 NodesSycamoreN12M14 上得到，然后逐行检查发现了这个bug

num_nodes             51
num_edges             99
ban_edges              0

先粗略记录一下。

这是print代码

'''calculate the multiple and avoid repeat'''
contract_dims = node_dims_arys[node_i0] + node_dims_arys[node_i1]  # 计算收缩后的node 的邻接张量的维度 以及来源
contract_bool = node_bool_arys[node_i0] | node_bool_arys[node_i1]  # 计算收缩后的node 由哪些原初node 合成
# assert contract_dims.shape == (num_nodes, )
# assert contract_bool.shape == (num_nodes, )

print(';;;', i, node_i0, node_i1)
print(node_dims_arys[node_i0].numpy().astype(int))
print(node_dims_arys[node_i1].numpy().astype(int))
print(contract_dims.numpy().astype(int))
print(contract_bool.numpy().astype(int))

这是print内容。可以看到，对已经收缩的节点竟然进行了不可能的收缩，并且产生了多余的乘法次数。

;;; 52 tensor(3) tensor(9)
[  0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0 352   0   0   0   0  96   0   0   0 128   0   0   0   0 320   0 128   0   0   0  64   0  64  64   0   0   0   0 192   0  64  64   0]
[  0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0 352   0   0   0   0  96   0   0   0 128   0   0   0   0 320   0 128   0   0   0  64   0  64  64   0   0   0   0 192   0  64  64   0]
[  0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0 704   0   0   0   0 192   0   0   0 256   0   0   0   0 640   0 256   0   0   0 128   0 128 128   0   0   0   0 384   0 128 128   0]
[1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 1 1 1 1 0 1 1 1 0 1 1 1 1 0 1 0 1 0 1 0 1 0 0 1 1 0 0 0 1 0 0 1]

[Yonv1943]

The following code fix this bug in Vanilla (single) mode.

The if_diff fix this bug.

First, choose a reasonable data type. It is safer that use float32 instead of int/long for node_dims_tens .

        node_dims_tens = th.stack([self.node_dims_ten.clone() for _ in range(num_envs)]).type(th.float32)
        node_bool_tens = th.stack([self.node_bool_ten.clone() for _ in range(num_envs)]).type(th.bool)

"""Vanilla (single)"""
for j in range(num_envs):
    edge_i = edge_is[j]
    node_dims_arys = node_dims_tens[j]
    node_bool_arys = node_bool_tens[j]

    '''find two nodes of an edge_i'''
    node_i0, node_i1 = th.where(edges_ary == edge_i)[0]  # 找出这条edge 两端的node
    # assert isinstance(node_i0.item(), int)
    # assert isinstance(node_i1.item(), int)

    '''whether node_i0 and node_i1 are different'''
    if_diff = th.logical_not(node_bool_arys[node_i0, node_i1])

    '''calculate the multiple and avoid repeat'''
    contract_dims = node_dims_arys[node_i0] + node_dims_arys[node_i1] * if_diff  # 计算收缩后的node 的邻接张量的维数以及来源
    contract_bool = node_bool_arys[node_i0] | node_bool_arys[node_i1]  # 计算收缩后的node 由哪些原初node 合成
    # assert contract_dims.shape == (num_nodes, )
    # assert contract_bool.shape == (num_nodes, )

    # 收缩掉的edge 只需要算一遍乘法。因此上面对 两次重复的指数求和后乘以0.5
    mult_pow_time = contract_dims.sum(dim=0) - (contract_dims * contract_bool).sum(dim=0) * 0.5
    mult_pow_timess[j, i] = mult_pow_time * if_diff

    '''adjust two list: node_dims_arys, node_bool_arys'''
    # 如果两个张量是一样的，那么 `contract_bool & if_diff` 就会全部变成 False，让下面这行代码不修改任何数值
    contract_dims[contract_bool & if_diff] = 0  # 把收缩掉的边的乘法数量赋值为2**0，接下来不再参与乘法次数的计算
    node_dims_tens[j, contract_bool] = contract_dims.repeat(1, 1)  # 根据 bool 将所有收缩后的节点都刷新成相同的信息
    node_bool_tens[j, contract_bool] = contract_bool.repeat(1, 1)  # 根据 bool 将所有收缩后的节点都刷新成相同的信息

Vectorized version

"""Vectorized"""
'''find two nodes of an edge_i'''
vec_edges_ary: TEN = edges_ary[None, :, :]
vec_edges_is: TEN = edge_is[:, None, None]
res = th.where(vec_edges_ary == vec_edges_is)[1]
res = res.reshape((num_envs, 2))
node_i0s, node_i1s = res[:, 0], res[:, 1]
# assert node_i0s.shape == (num_envs, )
# assert node_i1s.shape == (num_envs, )

'''whether node_i0 and node_i1 are different'''
if_diffs = th.logical_not(node_bool_tens[vec_env_is, node_i0s, node_i1s])

'''calculate the multiple and avoid repeat'''
contract_dimss = node_dims_tens[vec_env_is, node_i0s] + node_dims_tens[
    vec_env_is, node_i1s] * if_diffs.unsqueeze(1)
contract_bools = node_bool_tens[vec_env_is, node_i0s] | node_bool_tens[vec_env_is, node_i1s]
# assert contract_dimss.shape == (num_envs, num_nodes)
# assert contract_bools.shape == (num_envs, num_nodes)

mult_pow_times = contract_dimss.sum(dim=1) - (contract_dimss * contract_bools).sum(dim=1) * 0.5
# assert mult_pow_times.shape == (num_envs, )
mult_pow_timess[:, i] = mult_pow_times * if_diffs

'''adjust two list: node_dims_arys, node_bool_arys'''
for j in range(num_envs):  # 根据 bool 将所有收缩后的节点都刷新成相同的信息
    contract_dims = contract_dimss[j]
    contract_bool = contract_bools[j]

    contract_dims[contract_bool & if_diffs[j]] = 0  # 把收缩掉的边的乘法数量赋值为2**0，接下来不再参与乘法次数的计算
    node_dims_tens[j, contract_bool] = contract_dims.repeat(1, 1)
    node_bool_tens[j, contract_bool] = contract_bool.repeat(1, 1)

[Yonv1943]

修改后，我还使用换底公式 # max_tmp_power / th.log2(th.tensor((10, ), device=device)) # Change of Base Formula 优化了修复这个bug之后带来的 溢出问题

        # 计算这个乘法个数时，即便用 float64 也偶尔会过拟合，所以先除以 2**temp_power ，求log10 后再恢复它
        max_tmp_power = mult_pow_timess.max(dim=1)[0] - 960  # automatically set `max - 960`, 960 < the limit 1024,
        multiple_times = (2 ** (mult_pow_timess - max_tmp_power.unsqueeze(1))).sum(dim=1)
        multiple_times = multiple_times.log10() + max_tmp_power / th.log2(th.tensor((10,), device=device))
        # max_tmp_power / th.log2(th.tensor((10, ), device=device))  # Change of Base Formula