contractor 使用#
概述#
在本教程中，我们将演示如何利用不同类型的张量网络 contractor 进行电路仿真，以实现更好的时空消耗平衡。contractor 的定制是 TensorCircuit 库的主要亮点之一，因为更好的 contractor 可以更好地利用 TensorNetwok 仿真引擎的强大功能。
设置#
[1]:
import tensorcircuit as tc
import numpy as np
import cotengra as ctg
import opt_einsum as oem
K = tc.set_backend("tensorflow")
试验体系#
我们为两个电路提供张量网络，并测试这两个系统的收缩效率，第一个系统小，第二个系统大。
[2]:
# 获取小系统的量子态
def small_tn():
n = 10
d = 4
param = K.ones([2 * d, n])
c = tc.Circuit(n)
c = tc.templates.blocks.example_block(c, param, nlayers=d)
return c.state()
[3]:
# 获得对超大型系统的期望
def large_tn():
n = 60
d = 8
param = K.ones([2 * d, n])
c = tc.Circuit(n)
c = tc.templates.blocks.example_block(c, param, nlayers=d, is_split=True)
#
return c.expectation([tc.gates.z(), [n // 2]], reuse=False)
opt-einsum#
opt-einsum 提供了几个 contractor 优化器，并与 TensorNetwork 包一起提供。 由于 TensorCircuit 建立在 TensorNetwork 之上，我们可以使用这些简单的 contractor 优化器。 尽管对于任何中等系统，只有贪婪优化器有效，但其他优化器具有指数缩放并且在电路仿真场景中失效。我们总是为 contractor 系统设置contraction_info=True（默认为False），它将打印包括 contraction size、flops 和 writes 在内的收缩信息摘要。 有关这些指标的定义，另请参阅 cotengra 文档。
[4]:
# 如果我们什么都不设置，默认优化器是贪婪的，即:
tc.set_contractor("greedy", debug_level=2, contraction_info=True)
# 我们设置 debug_level=2 以不真正运行收缩计算
# 即通过设置debug_level>0，只有收缩信息和返回形状正确，结果错误
[4]:
functools.partial(<function custom at 0x7fd5a0a3d430>, optimizer=<function contraction_info_decorator.<locals>.new_algorithm at 0x7fd588e281f0>, memory_limit=None, debug_level=2)
[5]:
small_tn()
------ contraction cost summary ------
log10[FLOPs]: 5.132 log2[SIZE]: 11 log2[WRITE]: 13.083
[5]:
<tf.Tensor: shape=(1024,), dtype=complex64, numpy=
array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
dtype=complex64)>
[6]:
large_tn()
------ contraction cost summary ------
log10[FLOPs]: 17.766 log2[SIZE]: 44 log2[WRITE]: 49.636
[6]:
<tf.Tensor: shape=(), dtype=complex64, numpy=0j>
[7]:
# 我们可以在 opt-einsum 中使用更多花哨的 contractor，他们不在 TensorNetwork 中定义
# custom_stateful 用于具有一次性生命周期的路径求解器的收缩路径查找器
tc.set_contractor(
"custom_stateful",
optimizer=oem.RandomGreedy,
max_time=60,
max_repeats=128,
minimize="size",
debug_level=2,
contraction_info=True,
)
[7]:
functools.partial(<function custom_stateful at 0x7fd5a0a3d4c0>, optimizer=<class 'opt_einsum.path_random.RandomGreedy'>, opt_conf=None, contraction_info=True, debug_level=2, max_time=60, max_repeats=128, minimize='size')
[8]:
small_tn()
------ contraction cost summary ------
log10[FLOPs]: 4.925 log2[SIZE]: 10 log2[WRITE]: 12.531
[8]:
<tf.Tensor: shape=(1024,), dtype=complex64, numpy=
array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
dtype=complex64)>
[9]:
large_tn()
------ contraction cost summary ------
log10[FLOPs]: 11.199 log2[SIZE]: 26 log2[WRITE]: 28.183
[9]:
<tf.Tensor: shape=(), dtype=complex64, numpy=0j>
cotengra#
对于更高级的 contractor，我们向市场上的 sota contractor 优化器寻求帮助：cotengra
[10]:
opt = ctg.ReusableHyperOptimizer(
methods=["greedy", "kahypar"],
parallel=True,
minimize="write",
max_time=120,
max_repeats=1024,
progbar=True,
)
tc.set_contractor(
"custom", optimizer=opt, preprocessing=True, contraction_info=True, debug_level=2
)
## 有关 cotengra 优化器的更多设置，请参阅参考资料
## https://cotengra.readthedocs.io/en/latest/advanced.html
## preprocessing=True 将所有单量子比特门合并到相邻的双量子比特门
[10]:
functools.partial(<function custom at 0x7fd5a0a3d430>, optimizer=<function contraction_info_decorator.<locals>.new_algorithm at 0x7fd588e28ee0>, memory_limit=None, debug_level=2, preprocessing=True)
[11]:
small_tn()
log2[SIZE]: 10.00 log10[FLOPs]: 4.90: 100%|█████████████████████████████████████████| 1024/1024 [00:28<00:00, 35.45it/s]
------ contraction cost summary ------
log10[FLOPs]: 4.900 log2[SIZE]: 10 log2[WRITE]: 12.255
[11]:
<tf.Tensor: shape=(1024,), dtype=complex64, numpy=
array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
dtype=complex64)>
[12]:
large_tn()
log2[SIZE]: 20.00 log10[FLOPs]: 9.50: 4%|█▊ | 43/1024 [02:09<49:22, 3.02s/it]
------ contraction cost summary ------
log10[FLOPs]: 9.501 log2[SIZE]: 20 log2[WRITE]: 24.090
[12]:
<tf.Tensor: shape=(), dtype=complex64, numpy=0j>
我们也可以在 cotengra 找到路径之后应用 subtree reconf，这通常会进一步（并且通常会大大）为 contraction 改善 flops 和 writes。实际上, subtree reconf 后期处理通常比增加优化器的搜索时间或重复次数更重要。
[13]:
opt = ctg.ReusableHyperOptimizer(
minimize="combo",
max_repeats=1024,
max_time=240,
progbar=True,
)
def opt_reconf(inputs, output, size, **kws):
tree = opt.search(inputs, output, size)
tree_r = tree.subtree_reconfigure_forest(
progbar=True, num_trees=10, num_restarts=20, subtree_weight_what=("size",)
)
return tree_r.get_path()
tc.set_contractor(
"custom",
optimizer=opt_reconf,
contraction_info=True,
preprocessing=True,
debug_level=2,
)
[13]:
functools.partial(<function custom at 0x7fd5a0a3d430>, optimizer=<function contraction_info_decorator.<locals>.new_algorithm at 0x7fd58c832700>, memory_limit=None, debug_level=2, preprocessing=True)
[14]:
small_tn()
log2[SIZE]: 10.00 log10[FLOPs]: 4.87: 100%|█████████████████████████████████████████| 1024/1024 [02:29<00:00, 6.86it/s]
log2[SIZE]: 10.00 log10[FLOPs]: 4.86: 100%|█████████████████████████████████████████████| 20/20 [00:31<00:00, 1.57s/it]
------ contraction cost summary ------
log10[FLOPs]: 4.859 log2[SIZE]: 10 log2[WRITE]: 12.583
[14]:
<tf.Tensor: shape=(1024,), dtype=complex64, numpy=
array([0.+0.j, 0.+0.j, 0.+0.j, ..., 0.+0.j, 0.+0.j, 0.+0.j],
dtype=complex64)>
[15]:
large_tn()
log2[SIZE]: 21.00 log10[FLOPs]: 9.62: 9%|███▉ | 93/1024 [04:04<40:49, 2.63s/it]
log2[SIZE]: 17.00 log10[FLOPs]: 8.66: 100%|█████████████████████████████████████████████| 20/20 [03:08<00:00, 9.44s/it]
------ contraction cost summary ------
log10[FLOPs]: 8.657 log2[SIZE]: 17 log2[WRITE]: 25.035
[15]:
<tf.Tensor: shape=(), dtype=complex64, numpy=0j>
我们也可以直接提取张量网络用于电路或可观测计算，我们可以使用我们喜欢的任何方法进行收缩。此外，所有这些 contractor 或我们定制的外部收缩仍然可以与 jit，自动微分等兼容。具体来说，收缩路径求解器虽然需要一些时间成本，但由于 jit 基础设施只测量一次（注意，为了演示使用，我们在这里不使用K.jit装饰我们的收缩函数）。