Справочник API torch.func

Создано: 11 июня 2025 г. | Последнее обновление: 04 мая 2026 г.

Преобразования функций

`vmap`	vmap — это векторизующее отображение; `vmap(func)` возвращает новую функцию, которая отображает `func` по некоторому измерению входных данных.
`grad`	Оператор `grad` помогает вычислять градиенты `func` относительно входных данных, указанных в `argnums`.
`grad_and_value`	Возвращает функцию для вычисления кортежа из градиента и прямого, или прямого прохода, вычисления.
`vjp`	Возвращает кортеж, содержащий результаты применения `func` к `primals`, и функцию, которая при получении `cotangents` вычисляет якобиан `func` относительно `primals` в режиме обратного распространения, умноженный на `cotangents`. VJP означает произведение вектора на якобиан.
`jvp`	Возвращает кортеж, содержащий результат `func(*primals)` и «якобиан `func`, вычисленный в точке `primals`», умноженный на `tangents`. JVP означает произведение якобиана на вектор.
`linearize`	Возвращает значение `func` в точке `primals` и линейное приближение в `primals`.
`jacrev`	Вычисляет якобиан `func` относительно аргументов с индексом `argnum` с помощью автоматического дифференцирования в режиме обратного распространения
`jacfwd`	Вычисляет якобиан `func` относительно аргументов с индексом `argnum` с помощью автоматического дифференцирования в прямом режиме
`hessian`	Вычисляет матрицу Гессе `func` относительно аргументов с индексом `argnum` с помощью стратегии «прямой режим поверх обратного».
`functionalize`	functionalize — это преобразование, которое можно использовать для удаления мутаций (промежуточных) и алиасинга из функции, сохраняя при этом ее семантику.
`rearrange`	Нативная реализация `einops.rearrange` — удобного для чтения способа интеллектуальной перестановки элементов многомерных тензоров.

Утилиты для работы с модулями torch.nn

В общем случае можно применить преобразование к функции, вызывающей torch.nn.Module. Например, ниже показано, как вычислить якобиан функции, которая принимает три значения и возвращает три значения:

model = torch.nn.Linear(3, 3)

def f(x):
    return model(x)

x = torch.randn(3)
jacobian = jacrev(f)(x)
assert jacobian.shape == (3, 3)

Однако, если вы хотите, например, вычислить якобиан по параметрам модели, необходимо иметь возможность создать функцию, в которой параметры являются входными данными. Для этого предназначена функция functional_call(): она принимает nn.Module, преобразованный parameters и входные данные для прямого прохода модуля. Она возвращает значение, полученное при выполнении прямого прохода модуля с замененными параметрами.

Вот как можно вычислить якобиан по параметрам

model = torch.nn.Linear(3, 3)

def f(params, x):
    return torch.func.functional_call(model, params, x)

x = torch.randn(3)
jacobian = jacrev(f)(dict(model.named_parameters()), x)

`functional_call`	Выполняет функциональный вызов модуля, заменяя его параметры и буферы заданными значениями.
`stack_module_state`	Подготавливает список модулей torch.nn для ансамблирования с помощью `vmap()`.
`replace_all_batch_norm_modules_`	На месте обновляет `root`, устанавливая `running_mean` и `running_var` в None, а для любого модуля nn.BatchNorm в `root` устанавливая track_running_stats в False

Если вам нужна информация об исправлении модулей Batch Norm, воспользуйтесь приведенными здесь рекомендациями

Настройка Batch Norm

Утилиты отладки

debug_unwrap

Снимает обертку с тензора functorch (например, BatchedTensor или GradTrackingTensor), возвращая исходный тензор.

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/func.api.html