torch.fx.experimental.const_fold.split_const_subgraphs

torch.fx.experimental.const_fold.split_const_subgraphs(module, skip_folding_node_fn=None, device_for_folded_attrs='cpu', is_impure_node=None) [source]

Просматривает module в поисках узлов, у которых все входы атрибутов являются константами, отделяет их в отдельный подграф констант и возвращает FoldedGraphModule, который при первом запуске выполняет этот подграф констант, чтобы задать атрибуты модуля перед выполнением неконстантной части графа.

skip_folding_node_fn, если он передан, может вызываться для узлов, принадлежащих вложенным подграфам call_module, а не только для узлов верхнего уровня: узел call_module сворачивается атомарно, поэтому пропускается, если пропущен любой узел внутри его подграфа. Следовательно, предикаты должны работать только с узлом; предикат, которому необходимо разрешить target узла в подмодуль, должен использовать node.graph.owning_module, а не захваченную ссылку на модуль верхнего уровня.

is_impure_node, если он передан, передаётся в eliminate_dead_code, чтобы DCE сохранял узлы, которые вызывающий код считает имеющими побочные эффекты, помимо проверки Node.is_impure() по умолчанию (например, операции с вариантом out, записывающие данные в предварительно выделенный буфер через аргумент-ключевое слово out=, который не объявлен изменяемым в схеме).

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

FoldedGraphModule