Add converter torch aten::histc (#2796)

Fixes #2794.

---------

Signed-off-by: Xavier Dupré <xadupre@microsoft.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: 3 people

onnxscript/function_libs/torch_lib/ops/core.py

@@ -4608,12 +4608,53 @@ def aten_hinge_embedding_loss(
     raise NotImplementedError()
 
 
+@torch_op("aten::histc", trace_only=True)
 def aten_histc(
     self: TensorType, bins: int = 100, min: float = 0.0, max: float = 0.0
 ) -> TensorType:
     """histc(Tensor self, int bins=100, Scalar min=0, Scalar max=0) -> Tensor"""
+    if min == max:
+        # This ONNXScript implementation precomputes static bin edges and cannot
+        # faithfully reproduce torch.histc's dynamic behavior when min == max
+        # (including the default min=0, max=0, which infers the range from data).
+        raise NotImplementedError(
+            f"aten_histc with min == max ({min}) is not supported in this export path."
+        )
+    delta = (max - min) / (bins * 1.0)
+    values = [min + delta * i for i in range(bins + 1)]
 
-    raise NotImplementedError()
+    flat_self = op.Reshape(self, [-1])
+    computation_type = self.dtype
+
+    cond = op.And(
+        op.GreaterOrEqual(flat_self, op.CastLike([min], self)),
+        op.LessOrEqual(flat_self, op.CastLike([max], self)),
+    )
+
+    assert self.type.dtype not in {ir.DataType.INT32, ir.DataType.INT64}, (
+        f"torch.histc only works on float but {self.type.dtype=}"
+    )
+
+    cond = op.And(cond, op.Not(op.IsNaN(flat_self)))
+    # max is included.
+    dtype = self.type.dtype.numpy()
+    values = np.array(values, dtype=dtype)
+    values[-1] = np.nextafter(values[-1], np.array(np.inf, dtype=dtype), dtype=dtype)
+    typed_values = op.Constant(value=ir.tensor(values, dtype=self.type.dtype))
+
+    clipped = op.Where(cond, flat_self, op.CastLike([min - 1], self))
+    bins = op.Unsqueeze(typed_values, [1])
+
+    less = op.Cast(
+        op.Less(op.Unsqueeze(clipped, [0]), bins),
+        to=computation_type,
+    )
+    sums = op.ReduceSum(less, [1], keepdims=0)
+    res = op.Sub(
+        op.Slice(sums, [1], op.Shape(sums), [0]),
+        op.Slice(sums, [0], [-1], [0]),
+    )
+    return res
 
 
 def aten_histogramdd(

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -914,6 +914,51 @@ def forward(self, x):
         )
         _testing.assert_onnx_program(onnx_program)
 
+    def test_aten_histc_float(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.histc(x, 3, 0, 2)
+
+        model = Model()
+        onnx_program = torch.onnx.export(
+            Model(),
+            (torch.rand(10, 10, 10),),
+            dynamo=True,
+            verbose=False,
+            dynamic_shapes=({0: "batch"},),
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+        for k in range(101):
+            with self.subTest(k=k):
+                inputs = (torch.tensor([(k - 1) / 49.0], dtype=torch.float32),)
+                expected = model(*inputs)
+                got = onnx_program.call_reference({"x": inputs[0]})
+                torch.testing.assert_close(expected, got[0])
+
+    @unittest.skip("see https://github.com/pytorch/pytorch/issues/174668")
+    def test_aten_histc_float16(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.histc(x, 60, -10, 10)
+
+        model = Model()
+        onnx_program = torch.onnx.export(
+            Model(),
+            (torch.rand((10, 10, 10), dtype=torch.float16),),
+            dynamo=True,
+            verbose=False,
+            dynamic_shapes=({0: "batch"},),
+        )
+        _testing.assert_onnx_program(onnx_program)
+
+        for k in range(101):
+            with self.subTest(k=k):
+                inputs = (torch.tensor([(k - 1) / 49.0], dtype=torch.float16),)
+                expected = model(*inputs)
+                got = onnx_program.call_reference({"x": inputs[0]})
+                torch.testing.assert_close(expected, got[0])
+
 
 if __name__ == "__main__":
     unittest.main()