From d5290b80526e6e4dd79a9df87988d1137865a972 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Tue, 12 Aug 2025 19:59:32 +1000 Subject: [PATCH 1/6] added identity as alias for tensor_copy and defined No-Op for TensorFromScalar --- pytensor/tensor/basic.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 931c7009b3..ca7c73578b 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -663,6 +663,10 @@ def c_code_cache_version(self): tensor_from_scalar = TensorFromScalar() +@_vectorize_node.register(TensorFromScalar) +def vectorize_tensor_from_scalar(op, node, batch_x): + return identity(batch_x).owner + class ScalarFromTensor(COp): __props__ = () @@ -2046,6 +2050,8 @@ def register_transfer(fn): """Create a duplicate of `a` (with duplicated storage)""" tensor_copy = Elemwise(ps.identity) pprint.assign(tensor_copy, printing.IgnorePrinter()) +identity = tensor_copy +pprint.assign(identity, printing.IgnorePrinter()) class Default(Op): @@ -4603,6 +4609,7 @@ def ix_(*args): "matrix_transpose", "default", "tensor_copy", + "identity", "transfer", "alloc", "identity_like", From f20abb9ddce93e745d43a6af9a7dd0c8086cb4e7 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Tue, 12 Aug 2025 20:00:48 +1000 Subject: [PATCH 2/6] refactor: jacobian should use tensorize --- pytensor/tensor/optimize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytensor/tensor/optimize.py b/pytensor/tensor/optimize.py index 99a3d8b444..2810bc1ad1 100644 --- a/pytensor/tensor/optimize.py +++ b/pytensor/tensor/optimize.py @@ -560,7 +560,7 @@ def L_op(self, inputs, outputs, output_grads): implicit_f = grad(inner_fx, inner_x) df_dx, *df_dtheta_columns = jacobian( - implicit_f, [inner_x, *inner_args], disconnected_inputs="ignore" + implicit_f, [inner_x, *inner_args], disconnected_inputs="ignore", vectorize=True ) grad_wrt_args = implict_optimization_grads( df_dx=df_dx, From 30a71c4d2b43e7d434934e2e234a47c9bc172734 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 13 Aug 2025 09:36:30 +0000 Subject: [PATCH 3/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pytensor/tensor/basic.py | 1 + pytensor/tensor/optimize.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index ca7c73578b..2ec008a72d 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -663,6 +663,7 @@ def c_code_cache_version(self): tensor_from_scalar = TensorFromScalar() + @_vectorize_node.register(TensorFromScalar) def vectorize_tensor_from_scalar(op, node, batch_x): return identity(batch_x).owner diff --git a/pytensor/tensor/optimize.py b/pytensor/tensor/optimize.py index 2810bc1ad1..1c1127c894 100644 --- a/pytensor/tensor/optimize.py +++ b/pytensor/tensor/optimize.py @@ -560,7 +560,10 @@ def L_op(self, inputs, outputs, output_grads): implicit_f = grad(inner_fx, inner_x) df_dx, *df_dtheta_columns = jacobian( - implicit_f, [inner_x, *inner_args], disconnected_inputs="ignore", vectorize=True + implicit_f, + [inner_x, *inner_args], + disconnected_inputs="ignore", + vectorize=True, ) grad_wrt_args = implict_optimization_grads( df_dx=df_dx, From 8fe62c66f4748019eed73e93daa78f8d669041e1 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Fri, 15 Aug 2025 19:13:30 +1000 Subject: [PATCH 4/6] removed redundant pprint --- pytensor/tensor/basic.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 2ec008a72d..bf9638c473 100644 --- a/pytensor/tensor/basic.py +++ b/pytensor/tensor/basic.py @@ -2052,7 +2052,6 @@ def register_transfer(fn): tensor_copy = Elemwise(ps.identity) pprint.assign(tensor_copy, printing.IgnorePrinter()) identity = tensor_copy -pprint.assign(identity, printing.IgnorePrinter()) class Default(Op): From 35df4b20d3192b855175681e3191e9a3e921f691 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Fri, 15 Aug 2025 19:49:35 +1000 Subject: [PATCH 5/6] refactor: added vectorize=True to all jacobians --- pytensor/tensor/optimize.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pytensor/tensor/optimize.py b/pytensor/tensor/optimize.py index 1c1127c894..d90235d632 100644 --- a/pytensor/tensor/optimize.py +++ b/pytensor/tensor/optimize.py @@ -819,7 +819,7 @@ def __init__( self.fgraph = FunctionGraph([variables, *args], [equations]) if jac: - jac_wrt_x = jacobian(self.fgraph.outputs[0], self.fgraph.inputs[0]) + jac_wrt_x = jacobian(self.fgraph.outputs[0], self.fgraph.inputs[0], vectorize=True) self.fgraph.add_output(atleast_2d(jac_wrt_x)) self.jac = jac @@ -899,8 +899,8 @@ def L_op( inner_x, *inner_args = self.fgraph.inputs inner_fx = self.fgraph.outputs[0] - df_dx = jacobian(inner_fx, inner_x) if not self.jac else self.fgraph.outputs[1] - df_dtheta_columns = jacobian(inner_fx, inner_args, disconnected_inputs="ignore") + df_dx = jacobian(inner_fx, inner_x, vectorize=True) if not self.jac else self.fgraph.outputs[1] + df_dtheta_columns = jacobian(inner_fx, inner_args, disconnected_inputs="ignore", vectorize=True) grad_wrt_args = implict_optimization_grads( df_dx=df_dx, From ee37c24f29c4e7a39ab479048bc2f9017ec0d34e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 Aug 2025 10:07:43 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pytensor/tensor/optimize.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/pytensor/tensor/optimize.py b/pytensor/tensor/optimize.py index d90235d632..67ca71a210 100644 --- a/pytensor/tensor/optimize.py +++ b/pytensor/tensor/optimize.py @@ -819,7 +819,9 @@ def __init__( self.fgraph = FunctionGraph([variables, *args], [equations]) if jac: - jac_wrt_x = jacobian(self.fgraph.outputs[0], self.fgraph.inputs[0], vectorize=True) + jac_wrt_x = jacobian( + self.fgraph.outputs[0], self.fgraph.inputs[0], vectorize=True + ) self.fgraph.add_output(atleast_2d(jac_wrt_x)) self.jac = jac @@ -899,8 +901,14 @@ def L_op( inner_x, *inner_args = self.fgraph.inputs inner_fx = self.fgraph.outputs[0] - df_dx = jacobian(inner_fx, inner_x, vectorize=True) if not self.jac else self.fgraph.outputs[1] - df_dtheta_columns = jacobian(inner_fx, inner_args, disconnected_inputs="ignore", vectorize=True) + df_dx = ( + jacobian(inner_fx, inner_x, vectorize=True) + if not self.jac + else self.fgraph.outputs[1] + ) + df_dtheta_columns = jacobian( + inner_fx, inner_args, disconnected_inputs="ignore", vectorize=True + ) grad_wrt_args = implict_optimization_grads( df_dx=df_dx,