Комплексные числа

Создание комплексных тензоров

Поддерживаются два комплексных типа данных: torch.cfloat и torch.cdouble

>>> x = torch.randn(2,2, dtype=torch.cfloat)
>>> x
tensor([[-0.4621-0.0303j, -0.2438-0.5874j],
     [ 0.7706+0.1421j,  1.2110+0.1918j]])

Все фабричные функции, кроме torch.linspace(), torch.logspace() и torch.arange(), поддерживаются для комплексных тензоров.

Переход от старого представления

Пользователи, которые обходили отсутствие комплексных тензоров, используя действительные тензоры формы $(..., 2)$, могут легко перейти на комплексные тензоры в своем коде с помощью torch.view_as_complex() и torch.view_as_real(). Обратите внимание, что эти функции не выполняют копирование, а возвращают представление входного тензора.

>>> x = torch.randn(3, 2)
>>> x
tensor([[ 0.6125, -0.1681],
     [-0.3773,  1.3487],
     [-0.0861, -0.7981]])
>>> y = torch.view_as_complex(x)
>>> y
tensor([ 0.6125-0.1681j, -0.3773+1.3487j, -0.0861-0.7981j])
>>> torch.view_as_real(y)
tensor([[ 0.6125, -0.1681],
     [-0.3773,  1.3487],
     [-0.0861, -0.7981]])

Доступ к действительной и мнимой частям

Действительные и мнимые значения комплексного тензора можно получить с помощью real и imag.

>>> y.real
tensor([ 0.6125, -0.3773, -0.0861])
>>> y.imag
tensor([-0.1681,  1.3487, -0.7981])

>>> y.real.mul_(2)
tensor([ 1.2250, -0.7546, -0.1722])
>>> y
tensor([ 1.2250-0.1681j, -0.7546+1.3487j, -0.1722-0.7981j])
>>> y.real.stride()
(2,)

Аргумент и модуль

Аргумент и абсолютное значение комплексного тензора можно вычислить с помощью torch.angle() и torch.abs().

>>> x1=torch.tensor([3j, 4+4j])
>>> x1.abs()
tensor([3.0000, 5.6569])
>>> x1.angle()
tensor([1.5708, 0.7854])

Линейная алгебра

Многие операции линейной алгебры, такие как torch.matmul(), torch.linalg.svd(), torch.linalg.solve() и т. д., поддерживают комплексные числа. Если вы хотите запросить поддержку операции, которая пока не поддерживается, проверьте, создана ли уже соответствующая задача, а если нет, создайте ее.

Сериализация

Комплексные тензоры можно сериализовать, сохраняя данные в виде комплексных значений.

>>> torch.save(y, 'complex_tensor.pt')
>>> torch.load('complex_tensor.pt')
tensor([ 0.6125-0.1681j, -0.3773+1.3487j, -0.0861-0.7981j])

Автоматическое дифференцирование

PyTorch поддерживает автоматическое дифференцирование для комплексных тензоров. Вычисляемый градиент — это сопряженная производная Виртингера, противоположное направление которой в точности совпадает с направлением наискорейшего спуска, используемым в алгоритме градиентного спуска. Поэтому все существующие оптимизаторы можно использовать с комплексными параметрами без дополнительных изменений. Подробнее см. примечание Автоматическое дифференцирование для комплексных чисел.

Оптимизаторы

Семантически выполнение шага оптимизатора PyTorch с комплексными параметрами эквивалентно выполнению шага того же оптимизатора над представлением этих комплексных параметров в виде действительных чисел с помощью torch.view_as_real(). В частности:

>>> params = [torch.rand(2, 3, dtype=torch.complex64) for _ in range(5)]
>>> real_params = [torch.view_as_real(p) for p in params]

>>> complex_optim = torch.optim.AdamW(params)
>>> real_optim = torch.optim.AdamW(real_params)

real_optim и complex_optim будут вычислять одинаковые обновления параметров, хотя между двумя оптимизаторами могут возникать небольшие численные расхождения, как и между оптимизаторами foreach и forloop, а также capturable и стандартными оптимизаторами. Подробнее см. раздел о численной точности.

В частности, хотя можно считать, что наш оптимизатор обрабатывает комплексные тензоры так же, как если бы он отдельно оптимизировал их p.real и p.imag составляющие, детали реализации отличаются. Обратите внимание, что представление torch.view_as_real() преобразует комплексный тензор в действительный тензор формы $(..., 2)$, тогда как при разделении комплексного тензора на два тензора получаются два тензора размера $(...)$. Это различие не влияет на поэлементные оптимизаторы (например, AdamW), но приведет к небольшим расхождениям в оптимизаторах, выполняющих глобальные редукции (например, LBFGS). Сейчас у нас нет оптимизаторов, выполняющих редукции по каждому тензору, поэтому такое поведение пока не определено. Создайте задачу, если у вас есть сценарий использования, требующий точного определения этого поведения.

Следующие подсистемы поддерживаются не полностью:

• Квантование
• JIT
• Разреженные тензоры
• Распределенные вычисления

Если какая-либо из этих возможностей полезна для вашего сценария использования, проверьте, создана ли уже соответствующая задача, а если нет, создайте ее.