torch.nn.functional.scaled_grouped_mm

torch.nn.functional.scaled_grouped_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a=None, swizzle_b=None, bias=None, offs=None, output_dtype=torch.bfloat16, contraction_dim=(), use_fast_accum=False) [исходный код]

scaled_grouped_mm(mat_a, mat_b, scale_a, scale_recipe_a, scale_b, scale_recipe_b, swizzle_a, swizzle_b, bias, offs,

output_dtype, use_fast_accum)

Выполняет групповое умножение матриц с масштабированием: grouped_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) – необязательное смещение, добавляемое к выходным данным
• offs (Tensor | None) – необязательные смещения в исходных тензорах, задающие начальные индексы групп
• output_dtype (dtype | None) – dtype выходного тензора
• contraction_dim (list[int] | tuple[int, ...]) – задаёт измерения, соответствующие $K$ в матричном умножении.
• use_fast_accum (bool) – включает или отключает быстрое накопление на tensor core (только для GPU Hopper)

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

Tensor