Skip to content

Commit 36fdfcf

Browse files
authored
[Gluon][Tutorial] Use autolayout in a few places (#7646)
cc @manman-ren
1 parent 150c274 commit 36fdfcf

File tree

3 files changed

+4
-2
lines changed

3 files changed

+4
-2
lines changed

python/triton/experimental/gluon/language/_core.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -389,7 +389,7 @@ def convert_layout(value, layout, assert_trivial=False, _semantic=None):
389389

390390

391391
@builtin
392-
def full(shape, value, dtype, layout, _semantic=None):
392+
def full(shape, value, dtype, layout=None, _semantic=None):
393393
"""
394394
Create a tensor filled with a scalar value, with specified shape, dtype, and layout.
395395

python/triton/experimental/gluon/language/_semantic.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,8 @@ def splat(self, value, shape, layout):
158158

159159
def full(self, shape, value, dtype, layout):
160160
scalar = self.make_scalar(value, dtype)
161+
if layout is None:
162+
layout = AutoLayout()
161163
return self.splat(scalar, shape, layout)
162164

163165
def convert_layout(self, value, layout, assert_trivial=False):

python/tutorials/gluon/01-attention-forward.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -753,7 +753,7 @@ def _softmax_tile(tile_id: gl.constexpr, config, M, desc_o, STAGE: gl.constexpr,
753753
prog = scheduler.get_program(pid)
754754

755755
offs_m = prog.start_m * config.BLOCK_M
756-
offs_m += gl.arange(tile_id * config.SPLIT_M, (1 + tile_id) * config.SPLIT_M, qk_slice_dim1)
756+
offs_m += gl.arange(tile_id * config.SPLIT_M, (1 + tile_id) * config.SPLIT_M)
757757

758758
m_i = gl.full([config.SPLIT_M], -float("inf"), gl.float32, qk_slice_dim1)
759759
l_i0 = gl.full([config.SPLIT_M], 0.0, gl.float32, gl.SliceLayout(1, sum_layout))

0 commit comments

Comments
 (0)