torch.func.functionalize

torch.func.functionalize(func, *, remove='mutations') [исходный код]

functionalize — это трансформация, которую можно использовать для удаления мутаций (промежуточных) и алиасинга из функции с сохранением её семантики.

functionalize(func) возвращает новую функцию с той же семантикой, что и func, но без всех промежуточных мутаций. Каждая операция inplace, выполняемая над промежуточным тензором: intermediate.foo_() заменяется эквивалентной операцией вне места: intermediate_updated = intermediate.foo().

functionalize полезна для передачи программы pytorch бэкендам или компиляторам, которые не могут легко представлять операции мутации или алиасинга.

Параметры:

• func (Callable[..., Any]) – Функция Python, принимающая один или несколько аргументов.
• remove (str) – Необязательный строковый аргумент, принимающий значение ‘mutations’ или ‘mutations_and_views’. Если передано ‘mutations’, все мутирующие операторы будут заменены их немутирующими эквивалентами. Если передано ‘mutations_and_views’, дополнительно все операторы алиасинга будут заменены их эквивалентами без алиасинга. Значение по умолчанию: ‘mutations’.

Возвращает:

Возвращает новую «функционализированную» функцию. Она принимает те же входные данные, что и func, и ведёт себя так же, но все мутации (и, при необходимости, алиасинг), выполняемые над промежуточными тензорами внутри функции, будут удалены.

Тип возвращаемого значения:

Callable[…, Any]

functionalize также удаляет мутации (и представления), выполненные над входными данными функции. Однако для сохранения семантики functionalize «восстановит» мутации после завершения трансформации: определит, какие входные тензоры «должны были» быть изменены, и при необходимости скопирует новые данные обратно во входные тензоры.

Пример:

>>> import torch
>>> from torch.fx.experimental.proxy_tensor import make_fx
>>> from torch.func import functionalize
>>>
>>> # A function that uses mutations and views, but only on intermediate tensors.
>>> def f(a):
...     b = a + 1
...     c = b.view(-1)
...     c.add_(1)
...     return b
...
>>> inpt = torch.randn(2)
>>>
>>> out1 = f(inpt)
>>> out2 = functionalize(f)(inpt)
>>>
>>> # semantics are the same (outputs are equivalent)
>>> print(torch.allclose(out1, out2))
True
>>>
>>> f_traced = make_fx(f)(inpt)
>>> f_no_mutations_traced = make_fx(functionalize(f))(inpt)
>>> f_no_mutations_and_views_traced = make_fx(functionalize(f, remove='mutations_and_views'))(inpt)
>>>
>>> print(f_traced.code)



def forward(self, a_1):
    add = torch.ops.aten.add(a_1, 1);  a_1 = None
    view = torch.ops.aten.view(add, [-1])
    add_ = torch.ops.aten.add_(view, 1);  view = None
    return add

>>> print(f_no_mutations_traced.code)



def forward(self, a_1):
    add = torch.ops.aten.add(a_1, 1);  a_1 = None
    view = torch.ops.aten.view(add, [-1]);  add = None
    add_1 = torch.ops.aten.add(view, 1);  view = None
    view_1 = torch.ops.aten.view(add_1, [2]);  add_1 = None
    return view_1

>>> print(f_no_mutations_and_views_traced.code)



def forward(self, a_1):
    add = torch.ops.aten.add(a_1, 1);  a_1 = None
    view_copy = torch.ops.aten.view_copy(add, [-1]);  add = None
    add_1 = torch.ops.aten.add(view_copy, 1);  view_copy = None
    view_copy_1 = torch.ops.aten.view_copy(add_1, [2]);  add_1 = None
    return view_copy_1


>>> # A function that mutates its input tensor
>>> def f(a):
...     b = a.view(-1)
...     b.add_(1)
...     return a
...
>>> f_no_mutations_and_views_traced = make_fx(functionalize(f, remove='mutations_and_views'))(inpt)
>>> #
>>> # All mutations and views have been removed,
>>> # but there is an extra copy_ in the graph to correctly apply the mutation to the input
>>> # after the function has completed.
>>> print(f_no_mutations_and_views_traced.code)



def forward(self, a_1):
    view_copy = torch.ops.aten.view_copy(a_1, [-1])
    add = torch.ops.aten.add(view_copy, 1);  view_copy = None
    view_copy_1 = torch.ops.aten.view_copy(add, [2]);  add = None
    copy_ = torch.ops.aten.copy_(a_1, view_copy_1);  a_1 = None
    return view_copy_1

Следует отметить несколько «сценариев сбоя» functionalize:

1. Как и другие трансформации torch.func, functionalize() не работает с функциями, напрямую использующими .backward(). То же верно и для torch.autograd.grad. Если вам нужно использовать autograd, градиенты можно вычислить напрямую с помощью functionalize(grad(f)).
2. Как и другие трансформации torch.func, functionalize() не работает с глобальным состоянием. Если вызвать functionalize(f) для функции, выполняющей операции с представлениями / мутациями нелокального состояния, функционализация просто ничего не сделает и передаст вызовы представлений / мутаций непосредственно бэкенду. Обойти это ограничение можно, поместив создание любого нелокального состояния в более крупную функцию, к которой затем применяется functionalize.
3. resize_() имеет некоторые ограничения: functionalize работает только с программами, использующими resize_()`, если изменяемый размер тензор не является представлением.
4. as_strided() имеет некоторые ограничения: functionalize не работает с вызовами as_strided(), результатом которых являются тензоры с перекрывающейся памятью.

Наконец, для понимания функционализации полезно представлять, что большинство пользовательских программ pytorch написано с использованием публичного API torch. При выполнении операторы torch обычно декомпозируются во внутренний API C++ «ATen». Логика функционализации реализована целиком на уровне ATen. Функционализация умеет преобразовывать каждый оператор алиасинга в ATen в эквивалент без алиасинга (например, tensor.view({-1}) -> at::view_copy(tensor, {-1})), а каждый мутирующий оператор в ATen — в эквивалентный немутирующий оператор (например, tensor.add_(1) -> at::add(tensor, -1)), отслеживая алиасы и мутации отдельно, чтобы понимать, когда нужно восстановить исходное состояние. Информация о том, какие операторы ATen выполняют алиасинг или мутации, берётся из pytorch/pytorch.