softmax

Author: xadupre

_unittests/ut_reference/test_torch_evaluator.py

@@ -1,4 +1,5 @@
 import unittest
+from typing import Optional
 import numpy as np
 import onnx
 import onnx.helper as oh
@@ -22,7 +23,7 @@ def test_kernels(self):
         kernel = ker[key]
         self.assertEqual("Add_1", kernel.__name__)
 
-    def _finalize_test(self, model, *args):
+    def _finalize_test(self, model, *args, atol: Optional[float] = None):
         onnx.checker.check_model(model)
         feeds = dict(zip([i.name for i in model.graph.input], args))
 
@@ -31,7 +32,7 @@ def _finalize_test(self, model, *args):
         )
         rt = TorchEvaluator(model)
         got = rt.run(None, feeds)
-        self.assertEqualAny(expected, [g.detach().numpy() for g in got])
+        self.assertEqualAny(expected, [g.detach().numpy() for g in got], atol=atol)
 
     def test_op_binary(self):
         model = oh.make_model(
@@ -260,6 +261,21 @@ def test_op_gather(self):
             torch.tensor([0, 1, 3], dtype=torch.int64),
         )
 
+    def test_op_softmax(self):
+        model = oh.make_model(
+            oh.make_graph(
+                [oh.make_node("Softmax", ["X"], ["Z"], axis=0)],
+                "dummy",
+                [oh.make_tensor_value_info("X", TFLOAT, ["a", "b", "c"])],
+                [oh.make_tensor_value_info("Z", TFLOAT, ["a", "b", "c"])],
+            ),
+            ir_version=9,
+            opset_imports=[oh.make_opsetid("", 18)],
+        )
+        self._finalize_test(
+            model, torch.abs(torch.rand(3, 4, 5, dtype=torch.float32)), atol=1e-6
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -2,4 +2,5 @@
 from .access_ops import Gather_1, Slice_13
 from .binary_ops import Add_1, Div_1, MatMul_1, Mul_1, Sub_1
 from .other_ops import Cast_6, Concat_1, Transpose_1
+from .nn_ops import Softmax_13
 from .shape_ops import Reshape_5, Shape_15, Squeeze_13, Unsqueeze_13

onnx_diagnostic/reference/torch_ops/nn_ops.py

@@ -0,0 +1,22 @@
+from typing import Optional
+import onnx
+import torch
+from ...helpers.torch_helper import onnx_dtype_to_torch_dtype
+from . import OpRun, OpRunValue
+
+
+class Softmax_13(OpRun):
+    "Softmax"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.axis = self.get_attribute_int(node, "axis", -1)
+        assert isinstance(self.axis, int), f"Unexpected value for attribute axis={self.axis!r}"
+        # this is out of spec
+        stash_type = self.get_attribute_int(node, "stash_type", None)
+        self.stash_type = None if stash_type is None else onnx_dtype_to_torch_dtype(stash_type)
+
+    def run(self, data: OpRunValue) -> OpRunValue:
+        return OpRunValue(
+            torch.nn.functional.softmax(data.tensor, dim=self.axis, dtype=self.stash_type)
+        )