torch.cuda.jiterator._create_multi_output_jit_fn

torch.cuda.jiterator._create_multi_output_jit_fn(code_string, num_outputs, **kwargs) [исходный код]

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

Параметры:

• 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 выходных значений