Function _indicator

src/fused_mm_sampling/tl_argsort.py:107–110

(n_dims: core.constexpr, j: core.constexpr)

@triton.jit
def _indicator(n_dims: core.constexpr, j: core.constexpr):
    ar = core.arange(0, 2)
    ar = core.reshape(ar, [1] * (n_dims - j - 1) + [2] + [1] * j)
    return ar