torch.cuda.jiterator._create_multi_output_jit_fn

torch.cuda.jiterator._create_multi_output_jit_fn(code_string, num_outputs, **kwargs) [source]

Создать ядро CUDA, сгенерированное jiterator, для оператора поэлементного действия, поддерживающего возврат одного или нескольких результатов.

Параметры

• code_string (str) – Строка кода CUDA, подлежащая компиляции jiterator. Входной фунтор должен возвращать значение по ссылке.
• num_outputs (int) – количество результатов, возвращаемых ядром
• kwargs (Dict, optional) – Параметры ключевых слов для сгенерированной функции

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

Callable

Пример:

code_string = "template <typename T> void my_kernel(T x, T y, T alpha, T& out) { out = -x + alpha * y; }"
jitted_fn = create_jit_fn(code_string, alpha=1.0)
a = torch.rand(3, device='cuda')
b = torch.rand(3, device='cuda')
# invoke jitted function like a regular python function
result = jitted_fn(a, b, alpha=3.14)

Предупреждение

Этот API находится в стадии бета-тестирования и может быть изменён в будущих выпусках.

Предупреждение

Этот API поддерживает только до 8 входов и 8 выходов