Table of Contents
Shortcuts

# torch.cuda.jiterator._create_multi_output_jit_fn

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

  Create a jiterator-generated cuda kernel for an elementwise op that supports returning one or more outputs.

  Parameters:

    - code_string (str) – CUDA code string to be compiled by jiterator. The entry functor must return value by reference.

    - num_outputs (int) – number of outputs return by the kernel

    - kwargs (Dict, optional) – Keyword arguments for generated function

  Return type:

    Callable

  Example:

      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)

  Warning

  This API is in beta and may change in future releases.

  Warning

  This API only supports up to 8 inputs and 8 outputs

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.