triton.jit

triton.jit(*args, **kwargs)

Decorator for JIT-compiling a function using the Triton compiler.

Note

When a jit’d function is called, torch.tensor arguments are implicitly converted to pointers using the .data_ptr() method.

Note

This function will be compiled and run on the GPU. It will only have access to:

  • python primitives,

  • objects within the triton.language package,

  • arguments to this function,

  • other jit’d functions

Parameters

fn (Callable) – the function to be jit-compiled