torch.nn.functional.scaled_mm

torch.nn.functional.scaled_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a=None, swizzle_b=None, bias=None, output_dtype=torch.bfloat16, contraction_dim=(), use_fast_accum=False) [source]

scaled_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a, swizzle_b, bias, output_dtype,

contraction_dim, use_fast_accum)

Выполняет умножение матриц с масштабированием: mm(mat_a, mat_b), где масштабирование mat_a и mat_b задаётся соответственно параметрами scale_recipe_a и scale_recipe_b.

Параметры:

• scale_a (Tensor | list[Tensor]) – тензор, содержащий коэффициенты масштабирования для декодирования mat_a
• scale_recipe_a (_ScalingType | list[_ScalingType]) – перечисление, описывающее способ масштабирования mat_a
• scale_b (Tensor | list[Tensor]) – тензор, содержащий коэффициенты масштабирования для декодирования mat_b
• scale_recipe_b (_ScalingType | list[_ScalingType]) – перечисление, описывающее способ масштабирования mat_b
• swizzle_a (_SwizzleType | list[_SwizzleType] | None) – перечисление, описывающее схему перестановки (если она есть) для scale_a
• swizzle_b (_SwizzleType | list[_SwizzleType] | None) – перечисление, описывающее схему перестановки (если она есть) для scale_b
• bias (Tensor | None) – необязочное слагаемое смещения, добавляемое к результату
• output_dtype (dtype | None) – тип данных выходного тензора
• contraction_dim (list[int] | tuple[int, ...]) – описывает, какие размерности являются $K$ в операции умножения матриц.
• use_fast_accum (bool) – включает или отключает быстрое накопление на Tensor Core (только для GPU Hopper)

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

Tensor