frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix typo in doc and import for torch._library.triton (#144882)

Browse Source 
Previously, the doc's suggested `from torch._library.triton import wrap_triton, triton_op` doesn't work because wrap_triton is not imported in torch/_library/__init__.py but `from torch.library import wrap_triton` works. This PR imports wrap_triton and fix the doc.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/144882
Approved by: https://github.com/zou3519
This commit is contained in:
Yidi Wu
2025-01-16 16:45:11 -08:00
committed by PyTorch MergeBot
parent 18eba9575f
commit 567552b98b
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torch/_library/__init__.py
View File

@ -3,4 +3,4 @@ import torch._library.fake_impl
import torch._library.simple_registry
import torch._library.utils
from torch._library.fake_class_registry import register_fake_class
from torch._library.triton import capture_triton, triton_op
from torch._library.triton import capture_triton, triton_op, wrap_triton

2
torch/_library/triton.py
View File

@ -60,7 +60,7 @@ def triton_op(
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
>>> import torch
>>> from torch._library import triton_op, wrap_triton
>>> from torch.library import triton_op, wrap_triton
>>>
>>> import triton
>>> from triton import language as tl