PyTorch Softmax Ops

pytensor/link/pytorch/dispatch/elemwise.py

@@ -40,11 +40,14 @@ def dimshuffle(x):
 @pytorch_funcify.register(Softmax)
 def pytorch_funcify_Softmax(op, **kwargs):
     axis = op.axis
+    dtype = kwargs["node"].outputs[0].dtype
 
-    def softmax(x):
-        if not torch.is_floating_point(x):
-            x = x.to(torch.float32)
+    if not dtype.startswith("float"):
+        raise NotImplementedError(
+            "Pytorch Softmax is not currently implemented for non-float types."
+        )
 
+    def softmax(x):
         if axis is not None:
             return torch.softmax(x, dim=axis)
         else:
@@ -56,11 +59,14 @@ def softmax(x):
 @pytorch_funcify.register(LogSoftmax)
 def pytorch_funcify_LogSoftmax(op, **kwargs):
     axis = op.axis
+    dtype = kwargs["node"].outputs[0].dtype
 
-    def log_softmax(x):
-        if not torch.is_floating_point(x):
-            x = x.to(torch.float32)
+    if not dtype.startswith("float"):
+        raise NotImplementedError(
+            "Pytorch LogSoftmax is not currently implemented for non-float types."
+        )
 
+    def log_softmax(x):
         if axis is not None:
             return torch.log_softmax(x, dim=axis)
         else:

tests/link/pytorch/test_elemwise.py

@@ -57,24 +57,40 @@ def test_pytorch_elemwise():
     compare_pytorch_and_py(fg, [[0.9, 0.9]])
 
 
+@pytest.mark.parametrize("dtype", ["float64", "int64"])
 @pytest.mark.parametrize("axis", [None, 0, 1])
-def test_softmax(axis):
-    x = matrix("x")
+def test_softmax(axis, dtype):
+    x = matrix("x", dtype=dtype)
     out = softmax(x, axis=axis)
     fgraph = FunctionGraph([x], [out])
     test_input = np.arange(6, dtype=config.floatX).reshape(2, 3)
 
-    compare_pytorch_and_py(fgraph, [test_input])
+    if dtype == "int64":
+        with pytest.raises(
+            NotImplementedError,
+            match="Pytorch Softmax is not currently implemented for non-float types.",
+        ):
+            compare_pytorch_and_py(fgraph, [test_input])
+    else:
+        compare_pytorch_and_py(fgraph, [test_input])
 
 
+@pytest.mark.parametrize("dtype", ["float64", "int64"])
 @pytest.mark.parametrize("axis", [None, 0, 1])
-def test_logsoftmax(axis):
-    x = matrix("x")
+def test_logsoftmax(axis, dtype):
+    x = matrix("x", dtype=dtype)
     out = log_softmax(x, axis=axis)
     fgraph = FunctionGraph([x], [out])
     test_input = np.arange(6, dtype=config.floatX).reshape(2, 3)
 
-    compare_pytorch_and_py(fgraph, [test_input])
+    if dtype == "int64":
+        with pytest.raises(
+            NotImplementedError,
+            match="Pytorch LogSoftmax is not currently implemented for non-float types.",
+        ):
+            compare_pytorch_and_py(fgraph, [test_input])
+    else:
+        compare_pytorch_and_py(fgraph, [test_input])
 
 
 @pytest.mark.parametrize("axis", [None, 0, 1])