Preserve node output dtype when lifting scalars to attrs in XNNPACK quantizer

backends/xnnpack/quantizer/xnnpack_quantizer_utils.py

@@ -1157,15 +1157,15 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu
             prefix = "_tensor_constant_"
             get_new_attr_name = get_new_attr_name_with_prefix(prefix)
             tensor_constant_name = get_new_attr_name(model)
-            float_tensor = torch.tensor(float(args[i]))
-            model.register_buffer(tensor_constant_name, float_tensor)
+            scalar_tensor = torch.tensor(args[i], dtype=n.meta["val"].dtype)
+            model.register_buffer(tensor_constant_name, scalar_tensor)
             fake_mode = n.meta["val"].fake_mode
             with model.graph.inserting_before(n):
                 get_attr_node = model.graph.create_node(
                     "get_attr", tensor_constant_name, (), {}
                 )
                 get_attr_node.meta["val"] = fake_mode.from_tensor(
-                    float_tensor, static_shapes=True
+                    scalar_tensor, static_shapes=True
                 )
                 new_args.append(get_attr_node)
         n.args = tuple(new_args)

backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py

@@ -1122,6 +1122,30 @@ def forward(self, x):
             node_list,
         )
 
+    def test_int64_scalar_add_used_as_index(self):
+        """Scalars lifted to attrs must keep the op's output dtype; an int64
+        add chain used as an index must not be promoted to float32."""
+
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return x[:, torch.arange(4) + 0]
+
+        quantizer = XNNPACKQuantizer()
+        quantization_config = get_symmetric_quantization_config(is_per_channel=True)
+        quantizer.set_global(quantization_config)
+        example_inputs = (torch.randn(1, 4, 5),)
+        m = export(M(), example_inputs, strict=True).module()
+        m = quantizer.transform_for_annotation(m)
+        lifted_constants = [
+            m.get_buffer(n.target)
+            for n in m.graph.nodes
+            if n.op == "get_attr" and n.target.startswith("_tensor_constant_")
+        ]
+        self.assertEqual(len(lifted_constants), 1)
+        self.assertEqual(lifted_constants[0].dtype, torch.int64)
+        m = prepare_pt2e(m, quantizer)
+        m(*example_inputs)
+
     def test_cat_same_node(self):
         """Ensure that concatenating the same node does not cause any unexpected behavior"""