Add softplus

pymc-devs · Nov 2, 2024 · f1d7852 · f1d7852
1 parent a7a62f9
commit f1d7852
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pytensor/link/pytorch/dispatch/scalar.py b/pytensor/link/pytorch/dispatch/scalar.py
@@ -7,6 +7,7 @@
     Cast,
     ScalarOp,
 )
+from pytensor.scalar.math import Softplus
 
 
 @pytorch_funcify.register(ScalarOp)
@@ -56,3 +57,8 @@ def cast(x):
         return x.to(dtype=dtype)
 
     return cast
+
+
+@pytorch_funcify.register(Softplus)
+def pytorch_funcify_Softplus(op, node, **kwargs):
+    return torch.nn.Softplus()