gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] Added type check in semantic arange (#572)

Browse Source
This commit is contained in:
vesuppi
2022-07-03 15:25:37 -07:00
committed by GitHub
parent 4bf509889b
commit 22105bc33b
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/triton/language/semantic.py
View File

@@ -436,6 +436,9 @@ def not_equal(input: tl.tensor,
def arange(start: int, end: int, builder: ir.builder) -> tl.tensor: def arange(start: int, end: int, builder: ir.builder) -> tl.tensor:
if not isinstance(start, int) or not isinstance(end, int):
raise ValueError("arange's arguments must be of type tl.constexpr")
shape = [end - start] shape = [end - start]
ret_ty = tl.block_type(tl.int32, shape) ret_ty = tl.block_type(tl.int32, shape)
return tl.tensor(builder.get_range(start, end), ret_ty) return tl.tensor(builder.get_range(start, end), ret_ty)