Better docstring

Author: Ian Schweer

pytensor/link/pytorch/dispatch/elemwise.py

@@ -185,10 +185,9 @@ def softmax_grad(dy, sm):
 
 def elemwise_ravel_fn(base_fn, op, node, **kwargs):
     """
-    ScalarLoop + Elemwise is too common
-    to not work, but https://github.com/pymc-devs/pytensor/issues/1031,
-    vmap won't allow it. Instead, we can ravel all the inputs, broadcasted
-    according to torch
+    Dispatch methods using `.item()` (ScalarLoop + Elemwise) is common, but vmap
+    in torch has a limitation: https://github.com/pymc-devs/pytensor/issues/1031,
+    Instead, we can ravel all the inputs, broadcasted according to torch
     """
 
     n_outputs = len(node.outputs)

tests/link/pytorch/test_basic.py

@@ -393,6 +393,7 @@ def test_pytorch_softplus():
     f = FunctionGraph([x], [out])
     compare_pytorch_and_py(f, [np.random.rand(3)])
 
+
 def test_ScalarLoop():
     n_steps = int64("n_steps")
     x0 = float64("x0")