From 58e790bade586777a8e945b3b62bfb5f31081c55 Mon Sep 17 00:00:00 2001 From: Kazuaki Ishizaki Date: Tue, 25 Aug 2026 13:09:19 +0000 Subject: [PATCH] initial commit Signed-off-by: Kazuaki Ishizaki --- utils/model_ops/utils/torchop_yaml.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/utils/model_ops/utils/torchop_yaml.py b/utils/model_ops/utils/torchop_yaml.py index 487416dc..a8c73292 100644 --- a/utils/model_ops/utils/torchop_yaml.py +++ b/utils/model_ops/utils/torchop_yaml.py @@ -411,7 +411,7 @@ def _resolve_node_comments(comments): "torch.index_select", "torch.select_scatter", ] - _SPECIAL_INTLIMIT_OPS = _INDEX_INTLIMIT_OPS + ["torch.getitem"] + _SPECIAL_INTLIMIT_OPS = _INDEX_INTLIMIT_OPS + ["torch.getitem", "torch.setitem"] @staticmethod def _compute_randintlimit(op_name, i, dtype, saved_shape, san_args): @@ -419,11 +419,17 @@ def _compute_randintlimit(op_name, i, dtype, saved_shape, san_args): return 1000 if i == 0: return 1000 - if op_name == "torch.getitem" and "int" in str(dtype): - TorchOpCollector.log_function[TorchOpCollector.log_mthd]( - f"i: {i}, saved_shape: {saved_shape}, op_name: {op_name}, dtype: {dtype}, san_args: {san_args}" - ) - return saved_shape[0] + # getitem/setitem: ``a[idx]`` / ``a[idx] = v`` index dim 0 of arg 0, so the + # index tensor at i == 1 is bounded by saved_shape[0]. Guard on i to avoid + # bounding setitem's value tensor (i == 2), which carries no index + # semantics even when its dtype is integral. + if op_name in ("torch.getitem", "torch.setitem"): + if i == 1 and "int" in str(dtype) and saved_shape: + TorchOpCollector.log_function[TorchOpCollector.log_mthd]( + f"i: {i}, saved_shape: {saved_shape}, op_name: {op_name}, dtype: {dtype}, san_args: {san_args}" + ) + return saved_shape[0] + return 1000 if op_name in TorchOpCollector._INDEX_INTLIMIT_OPS: dim_index = 2 if op_name == "torch.select_scatter" else 1 if i == dim_index + 1 and "int" in str(dtype):