cuquantum.contract

cuquantum.contract(subscripts, *operands, options=None, optimize=None, stream=None, return_info=False)[源代码]

在操作数上评估爱因斯坦求和约定。

爱因斯坦求和表达式支持显式和隐式形式。除了下标格式外，交错格式也支持作为指定操作数及其模式标签的一种方式。有关操作数类型和示例的更多详细信息，请参阅 Network。

参数

  • subscripts – 模式标签（下标），定义爱因斯坦求和表达式为逗号分隔的字符序列。表达式中允许使用 Unicode 字符，从而扩展了可以使用爱因斯坦求和约定指定的张量网络的规模。

  • operands – 张量（类 ndarray 对象）序列。当前支持的类型为 numpy.ndarray、cupy.ndarray 和 torch.Tensor。

  • qualifiers – 将张量限定符指定为长度等于操作数数量的 numpy.ndarray tensor_qualifiers_dtype 对象。

  • options – 将张量网络的选项指定为 NetworkOptions 对象。或者，也可以提供包含 NetworkOptions 构造函数参数的 dict。如果未指定，该值将设置为默认构造的 NetworkOptions 对象。

  • optimize – 此参数将路径优化选项指定为 OptimizerOptions 对象。或者，也可以提供包含 OptimizerOptions 构造函数参数的字典。如果未指定，该值将设置为默认构造的 OptimizerOptions 对象。

  • stream – 提供 CUDA 流以用于自动调优操作。可接受的输入包括 cudaStream_t（作为 Python int）、cupy.cuda.Stream 和 torch.cuda.Stream。如果未提供流，将使用当前流。

  • return_info – 如果为 true，则还将返回有关最佳收缩顺序的信息。

返回

如果 return_info 为 False，则输出张量（类 ndarray 对象），其类型和设备与包含收缩结果的操作数相同；否则，返回一个 2 元组，其中包含输出张量和一个 OptimizerInfo 对象，该对象包含有关最佳收缩顺序等信息。

注意

鼓励用户自行维护库句柄，以减少上下文初始化时间

from cuquantum import cutensornet as cutn
from cuquantum import contract, NetworkOptions

handle = cutn.create()
network_opts = NetworkOptions(handle=handle, ...)
out = contract(..., options=network_opts, ...)
# ... the same handle can be reused for further calls ...
# when it's done, remember to destroy the handle
cutn.destroy(handle)

示例

使用 NumPy 操作数

>>> from cuquantum import contract
>>> import numpy as np
>>> a = np.arange(6.).reshape(3, 2)
>>> b = np.arange(6.).reshape(2, 3)

使用显式形式执行矩阵乘法。结果 r 是 NumPy ndarray（计算在 GPU 上执行）

>>> r = contract('ij,jk->ik', a, b)

隐式形式

>>> r = contract('ij,jk', a, b)

使用字符作为模式标签的交错格式

>>> r = contract(a, ['i', 'j'], b, ['j', 'k'], ['i', 'k'], return_info=True)

使用字符串标签作为模式标签和隐式形式的交错格式

>>> r = contract(a, ['first', 'second'], b, ['second', 'third'])

使用整数模式标签和显式形式的交错格式

>>> r = contract(a, [1, 2], b, [2, 3], [1, 3])

获取有关最佳收缩路径的信息 i 以及结果 r

>>> r, i = contract('ij,jk', a, b, return_info=True)

为张量网络提供选项

>>> from cuquantum import NetworkOptions
>>> n = NetworkOptions(device_id=1)
>>> r = contract('ij,jk->ik', a, b, options=n)

或者，选项可以作为 dict 而不是 NetworkOptions 对象提供

>>> r = contract('ij,jk->ik', a, b, options={'device_id': 1})

指定优化器的选项

>>> from cuquantum import OptimizerOptions, PathFinderOptions
>>> p = PathFinderOptions(imbalance_factor=230, cutoff_size=8)
>>> o = OptimizerOptions(path=p, seed=123)
>>> r = contract('ij,jk,kl', a, b, a, optimize=o)

或者，上面的选项可以作为 dict 提供

>>> r = contract('ij,jk,kl', a, b, a, optimize={'path': {'imbalance_factor': 230, 'cutoff_size': 8}, 'seed': 123})

直接指定路径

>>> o = OptimizerOptions(path = [(0,2), (0,1)])
>>> r = contract('ij,jk,kl', a, b, a, optimize=o)

使用省略号简写符号执行元素级乘法 \(a \odot b^T\)

>>> r = contract('...,...', a, b.T)

使用省略号简写符号获取双内积 \(a : b^T\)（实值张量的 Frobenius 内积）

>>> r = contract('...,...->', a, b.T)

使用 CuPy 操作数。结果 r 是与操作数位于同一设备上的 CuPy ndarray，dev 是您希望用于存储张量和计算收缩的系统上的任何有效设备 ID

>>> import cupy
>>> dev = 0
>>> with cupy.cuda.Device(dev):
...     a = cupy.arange(6.).reshape(3, 2)
...     b = cupy.arange(6.).reshape(2, 3)
>>> r = contract('ij,jk', a, b)

对于 PyTorch 操作数，此函数像本机 PyTorch 运算符一样工作，可以由 autograd 图跟踪，从而实现反向传播。结果 r 是与操作数位于同一设备 (dev) 上的 PyTorch 张量。要启用梯度计算，只需像往常一样将目标操作数的 requires_grad 属性设置为 True。如果显式传递了 stream，则用户必须按照 PyTorch CUDA 语义中概述的要求建立流排序。

>>> import torch
>>> dev = 0
>>> a = torch.arange(6., device=f'cuda:{dev}').reshape(3, 2)
>>> a.requires_grad_(True)
>>> b = torch.arange(6., device=f'cuda:{dev}').reshape(2, 3)
>>> b.requires_grad_(True)
>>> r = contract('ij,jk', a, b)
>>> r.backward(torch.ones_like(r))  # gradient w.r.t self is 1
>>> a.grad
tensor([[ 3., 12.],
        [ 3., 12.],
        [ 3., 12.]], device='cuda:0')
>>> b.grad
tensor([[6., 6., 6.],
        [9., 9., 9.]], device='cuda:0')