Support Tensor for clamp

test/modules/op/clamp.py

@@ -103,6 +103,21 @@ def get_example_inputs(self):
         return (torch.randn(5, 3) * 10,), {}
 
 
+class SimpleClampTensorMinMax(TestModuleBase):
+    """Clamp with scalar tensor buffers as used by Gemma4ClippableLinear."""
+
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("min_value", torch.tensor(-11.0))
+        self.register_buffer("max_value", torch.tensor(11.0))
+
+    def forward(self, x):
+        return torch.clamp(x, self.min_value, self.max_value)
+
+    def get_example_inputs(self):
+        return (torch.randn(5, 3) * 20,), {}
+
+
 class ClampIntInputFloatMinMax(TestModuleBase):
     def __init__(self):
         super().__init__()

test/unit_test/passes/test_cast_clamp_mixed_type_args.py

@@ -106,3 +106,29 @@ def test_pass(self):
         self.run_value_test(CastClampMixedTypeArgs())
         self.assertEqual(num_of_ops(self.exported_program(), ops.aten.clamp), 1)
         self.assertEqual(num_of_ops(self.exported_program(), ops.aten._to_copy), 1)
+
+
+class CastClampTensorBounds(torch.nn.Module):
+    """Use lifted scalar tensor buffers as Clamp min/max arguments."""
+
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("min_value", torch.tensor(-10.0))
+        self.register_buffer("max_value", torch.tensor(10.0))
+
+    def forward(self, x):
+        return torch.clamp(x, self.min_value, self.max_value)
+
+    def get_example_inputs(self):
+        return (torch.randn(5, 3) * 20,), {}
+
+
+class CastClampTensorBoundsTest(SinglePassValueTest):
+    def test_pass(self):
+        self.setup(CastClampTensorBounds())
+        self.assertEqual(num_of_ops(self.exported_program(), ops.aten.clamp), 1)
+
+        self.run_value_test(CastClampMixedTypeArgs())
+
+        self.assertEqual(num_of_ops(self.exported_program(), ops.aten.clamp), 1)
+        self.assertEqual(num_of_ops(self.exported_program(), ops.aten._to_copy), 0)

tico/passes/cast_clamp_mixed_type_args.py

@@ -100,6 +100,7 @@ def convert(self, exported_program: ExportedProgram, node: torch.fx.Node) -> boo
         graph = graph_module.graph
 
         # clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor
+        # clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor
         args = ClampArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
 
         input = args.input
@@ -113,20 +114,42 @@ def _convert_arg(arg, arg_name: str):
             if arg is None:
                 return False
 
-            arg_dtype = torch.tensor(arg).dtype
             arg_idx = node.args.index(arg)
-            if arg_dtype != output_dtype:
-                assert output_dtype in [torch.float, torch.int]
-                if output_dtype == torch.float:
-                    arg = float(arg)
-                else:
-                    arg = int(arg)
-                node.update_arg(arg_idx, arg)
+
+            if isinstance(arg, torch.fx.Node):
+                arg_dtype = extract_torch_dtype(arg)
+                if arg_dtype == output_dtype:
+                    return False
+
+                with graph.inserting_after(arg):
+                    to_copy = create_node(
+                        graph,
+                        torch.ops.aten._to_copy.default,
+                        (arg,),
+                        {"dtype": output_dtype},
+                        origin=arg,
+                    )
+                    set_new_meta_val(to_copy)
+                    node.update_arg(arg_idx, to_copy)
+
                 logger.debug(
-                    f"Casting {arg_name} value from {arg_dtype} to {output_dtype} for clamp operation at {node.name}"
+                    f"Inserting cast for {arg_name} from {arg_dtype} to "
+                    f"{output_dtype} for clamp operation at {node.name}"
                 )
                 return True
-            return False
+
+            arg_dtype = torch.tensor(arg).dtype
+            if arg_dtype == output_dtype:
+                return False
+
+            assert output_dtype in [torch.float, torch.int]
+            converted_arg = float(arg) if output_dtype == torch.float else int(arg)
+            node.update_arg(arg_idx, converted_arg)
+            logger.debug(
+                f"Casting {arg_name} value from {arg_dtype} to {output_dtype} "
+                f"for clamp operation at {node.name}"
+            )
+            return True
 
         modified |= _convert_arg(min, "min")
         modified |= _convert_arg(max, "max")

tico/utils/validate_args_kwargs.py

@@ -208,11 +208,12 @@ class CircleRMSNormArgs:
 class ClampArgs:
     """
     clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor
+    clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor
     """
 
     input: torch.fx.Node
-    min: Optional[Union[int, float]] = None
-    max: Optional[Union[int, float]] = None
+    min: Optional[Union[torch.fx.Node, int, float]] = None
+    max: Optional[Union[torch.fx.Node, int, float]] = None
 
 
 @enforce_type