Update vectorize_graph docstring

Docstring still uses old name of `vectorize`
ricardoV94 · Dec 15, 2023 · 60a9566 · 60a9566
1 parent 8ae14c2
commit 60a9566
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pytensor/graph/replace.py b/pytensor/graph/replace.py
@@ -256,15 +256,15 @@ def vectorize_graph(
         import pytensor
         import pytensor.tensor as pt
 
-        from pytensor.graph import vectorize
+        from pytensor.graph import vectorize_graph
 
         # Original graph
         x = pt.vector("x")
         y = pt.exp(x) / pt.sum(pt.exp(x))
 
         # Vectorized graph
         new_x = pt.matrix("new_x")
-        new_y = vectorize(y, replace={x: new_x})
+        new_y = vectorize_graph(y, replace={x: new_x})
 
         fn = pytensor.function([new_x], new_y)
         fn([[0, 1, 2], [2, 1, 0]])
@@ -277,7 +277,7 @@ def vectorize_graph(
         import pytensor
         import pytensor.tensor as pt
 
-        from pytensor.graph import vectorize
+        from pytensor.graph import vectorize_graph
 
         # Original graph
         x = pt.vector("x")
@@ -286,7 +286,7 @@ def vectorize_graph(
 
         # Vectorized graph
         new_x = pt.matrix("new_x")
-        [new_y1, new_y2] = vectorize([y1, y2], replace={x: new_x})
+        [new_y1, new_y2] = vectorize_graph([y1, y2], replace={x: new_x})
 
         fn = pytensor.function([new_x], [new_y1, new_y2])
         fn([[-10, 0, 10], [-11, 0, 11]])