torch.autograd.Function.vmap

static Function.vmap(info, in_dims, *args) [исходный код]

Определяет поведение этого autograd.Function внутри torch.vmap().

Чтобы torch.autograd.Function() поддерживал torch.vmap(), необходимо переопределить этот статический метод или присвоить generate_vmap_rule значение True (нельзя делать и то и другое).

Если вы решили переопределить этот статический метод, он должен принимать

• объект info в качестве первого аргумента. info.batch_size задаёт размер измерения, по которому выполняется vmap, а info.randomness — параметр случайности, передаваемый в torch.vmap().
• кортеж in_dims в качестве второго аргумента. Для каждого аргумента в args в in_dims имеется соответствующий Optional[int]. Это None, если аргумент не является тензором или если для него не выполняется vmap; в противном случае это целое число, указывающее, по какому измерению тензора выполняется vmap.
• *args, совпадающий с аргументами forward().

Статический метод vmap возвращает кортеж из (output, out_dims). Как и in_dims, out_dims должен иметь ту же структуру, что и output, и содержать по одному out_dim для каждого результата, указывающему, содержит ли результат измерение vmap и по какому индексу оно расположено.

Подробнее см. в разделе Расширение torch.func с помощью autograd.Function.