Skip to content

Commit 488d058

Browse files
Michal-Novomestskyjessegrabowski
authored andcommitted
refactor: added vectorize=True to all jacobians
1 parent 94c3d6d commit 488d058

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

pytensor/tensor/optimize.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -820,7 +820,7 @@ def __init__(
820820
self.fgraph = FunctionGraph([variables, *args], [equations])
821821

822822
if jac:
823-
jac_wrt_x = jacobian(self.fgraph.outputs[0], self.fgraph.inputs[0])
823+
jac_wrt_x = jacobian(self.fgraph.outputs[0], self.fgraph.inputs[0], vectorize=True)
824824
self.fgraph.add_output(atleast_2d(jac_wrt_x))
825825

826826
self.jac = jac
@@ -900,8 +900,8 @@ def L_op(
900900
inner_x, *inner_args = self.fgraph.inputs
901901
inner_fx = self.fgraph.outputs[0]
902902

903-
df_dx = jacobian(inner_fx, inner_x) if not self.jac else self.fgraph.outputs[1]
904-
df_dtheta_columns = jacobian(inner_fx, inner_args, disconnected_inputs="ignore")
903+
df_dx = jacobian(inner_fx, inner_x, vectorize=True) if not self.jac else self.fgraph.outputs[1]
904+
df_dtheta_columns = jacobian(inner_fx, inner_args, disconnected_inputs="ignore", vectorize=True)
905905

906906
grad_wrt_args = implict_optimization_grads(
907907
df_dx=df_dx,

0 commit comments

Comments
 (0)