tensorplay.library.triton_op

tensorplay.library.triton_op(name: str, fn: Callable[[...], Any] | None = None, /, *, mutates_args: Sequence[str] = (), device_types: Any = None, schema: str | None = None) → Callable[[Callable[[...], Any]], CustomOpDef][source]

The registered kernel(s) must launch their Triton kernels through wrap_triton() and only mutate arguments listed in mutates_args. Under tensorplay.compile the whole operator is captured as a single opaque node — the compiler never traces into the Triton launches. device_types is a