Add function for input preprocessing in numerical comparator (#11739)

Summary:
Pull Request resolved: #11739

This PR adds the process_input_to_tensor function to convert inputs to torch.Tensor on CPU with torch.float32 dtype for preprocessing inputs used later in the numerical comparator.

Differential Revision: D76745314

Author: Juntian Liu

devtools/inspector/_inspector_utils.py

@@ -690,3 +690,38 @@ def map_runtime_aot_intermediate_outputs(
             )
 
     return aot_runtime_mapping
+
+
+def process_input_to_tensor(input_data: Any) -> torch.Tensor:
+    """
+    Convert input_data into a torch.Tensor on CPU with dtype torch.float32.
+    This function handles the following types of input:
+    - Scalar (int or float): Converts to a 1D tensor with a single element.
+    - Tensor: Converts to a float32 tensor on CPU.
+    - List of Tensors: Stacks the tensors into a single float32 tensor on CPU.
+    The resulting tensor is detached, moved to CPU, and cast to torch.float32.
+    Parameters:
+    input_data (Any): The input data to be converted to a tensor. It can be a scalar,
+                      a tensor, or a list of tensors.
+    Returns:
+    torch.Tensor: A tensor on CPU with dtype torch.float32.
+    Raises:
+    ValueError: If the input_data cannot be converted to a tensor.
+    """
+    try:
+        # Check if the input is a scalar
+        if isinstance(input_data, (int, float)):
+            input_tensor = torch.tensor([input_data], dtype=torch.float32)
+        # Check if the input is a list of tensors
+        elif isinstance(input_data, list) and all(
+            isinstance(i, torch.Tensor) for i in input_data
+        ):
+            input_tensor = torch.stack(input_data).to(torch.float32)
+        else:
+            input_tensor = torch.as_tensor(input_data, dtype=torch.float32)
+    except Exception as e:
+        raise ValueError(
+            f"Cannot convert value of type {type(input_data)} to a tensor: {e}"
+        )
+    input_tensor = input_tensor.detach().cpu().float()
+    return input_tensor

devtools/inspector/tests/inspector_utils_test.py

@@ -36,6 +36,7 @@
     is_inference_output_equal,
     map_runtime_aot_intermediate_outputs,
     merge_overlapping_debug_handles,
+    process_input_to_tensor,
     TimeScale,
 )
 
@@ -317,6 +318,43 @@ def test_map_runtime_aot_intermediate_outputs_complex_chain(self):
         expected = {((1, 2, 3, 4, 5, 6), 300): ((2, 3, 4, 5, 6, 7), 350)}
         self.assertEqual(actual, expected)
 
+    def test_process_input_convertible_inputs(self):
+        # Scalar -> tensor
+        actual_output1 = process_input_to_tensor(5)
+        self.assertIsInstance(actual_output1, torch.Tensor)
+        self.assertEqual(actual_output1.dtype, torch.float32)
+        self.assertEqual(tuple(actual_output1.shape), (1,))
+        self.assertTrue(torch.allclose(actual_output1, torch.tensor([5.0])))
+        self.assertEqual(actual_output1.device.type, "cpu")
+
+        # Tensor of ints -> float32 CPU
+        t_int = torch.tensor([4, 5, 6], dtype=torch.int32)
+        actual_output2 = process_input_to_tensor(t_int)
+        self.assertIsInstance(actual_output2, torch.Tensor)
+        self.assertEqual(actual_output2.dtype, torch.float32)
+        self.assertTrue(torch.allclose(actual_output2, torch.tensor([4.0, 5.0, 6.0])))
+        self.assertEqual(actual_output2.device.type, "cpu")
+
+        # List of tensors -> stacked tensor float32 CPU
+        t_list = [torch.tensor([1]), torch.tensor([2]), torch.tensor([3])]
+        actual_output3 = process_input_to_tensor(t_list)
+        self.assertIsInstance(actual_output3, torch.Tensor)
+        self.assertEqual(actual_output3.dtype, torch.float32)
+        self.assertEqual(tuple(actual_output3.shape), (3, 1))
+        self.assertTrue(
+            torch.allclose(actual_output3, torch.tensor([[1.0], [2.0], [3.0]]))
+        )
+        self.assertEqual(actual_output3.device.type, "cpu")
+
+    def test_process_input_tensor_non_convertible_raises(self):
+        class X:
+            pass
+
+        with self.assertRaises(ValueError) as cm:
+            process_input_to_tensor(X())
+        msg = str(cm.exception)
+        self.assertIn("Cannot convert value of type", msg)
+
 
 def gen_mock_operator_graph_with_expected_map() -> (
     Tuple[OperatorGraph, Dict[int, OperatorNode]]