diff --git a/mindtorch/_apis/cpu.py b/mindtorch/_apis/cpu.py index 1da15e7a7..2650e81d4 100644 --- a/mindtorch/_apis/cpu.py +++ b/mindtorch/_apis/cpu.py @@ -1236,4 +1236,7 @@ def search_sorted(sorted_sequence, values, sorter, dtype, right): return legacy.search_sorted(sorted_sequence, values, sorter, dtype, right) def scatter_nd_update(input, indices, updates): - return legacy.scatter_nd_update(input, indices, updates, True) \ No newline at end of file + return legacy.scatter_nd_update(input, indices, updates, True) + +def triu_indices(row, col, offset, dtype): + return legacy.triu_indices(row, col, offset, dtype) diff --git a/mindtorch/_apis/gpu.py b/mindtorch/_apis/gpu.py index 815795018..52733f07e 100644 --- a/mindtorch/_apis/gpu.py +++ b/mindtorch/_apis/gpu.py @@ -1245,3 +1245,5 @@ def fft(input, n=None, dim=-1, norm="backward"): input = narrow(input, dim, 0, n) return legacy.fft_with_size(input, input.ndim, False, False, norm, True, ()) +def triu_indices(row, col, offset, dtype): + return legacy.triu_indices(row, col, offset, dtype) diff --git a/mindtorch/_apis/npu.py b/mindtorch/_apis/npu.py index 540462fca..4edec5138 100644 --- a/mindtorch/_apis/npu.py +++ b/mindtorch/_apis/npu.py @@ -1631,3 +1631,6 @@ def repeat_interleave_int(input, repeats, dim, output_size): def repeat_interleave_tensor(input, repeats, dim, output_size): return pyboost.repeat_interleave_tensor_op(input, repeats, dim, output_size) + +def triu_indices(row, col, offset, dtype): + return legacy.triu_indices(row, col, offset, dtype) diff --git a/mindtorch/ops/other.py b/mindtorch/ops/other.py index dfcd722de..1b76f9b7d 100644 --- a/mindtorch/ops/other.py +++ b/mindtorch/ops/other.py @@ -939,6 +939,8 @@ def triu(input, diagonal=0): return execute('triu', input, diagonal) # triu_indices +def triu_indices(row, col, offset=0, *, dtype=mindtorch.long, device='cpu', layout=mindtorch.strided): + return execute('triu_indices', row, col, offset, dtype, device=device) # unflatten @@ -1145,6 +1147,7 @@ def cosine_similarity(x1, x2, dim=1, eps=1e-8): "stop_gradient", "tril", "triu", + "triu_indices", "unflatten", "unfold", "contiguous",