Correct bad imports in optimize.py and expose it via pytensor.tensor.__init__

doc/library/tensor/index.rst

@@ -31,3 +31,4 @@ symbolic expressions using calls that look just like numpy calls, such as
     math_opt
     basic_opt
     functional
+    optimize

doc/library/tensor/optimize.rst

@@ -0,0 +1,11 @@
+========================================================
+:mod:`tensor.optimize` -- Symbolic Optimization Routines
+========================================================
+
+.. module:: tensor.conv
+   :platform: Unix, Windows
+   :synopsis: Symbolic Optimization Routines
+.. moduleauthor:: LISA, PyMC Developers, PyTensor Developers
+
+.. automodule:: pytensor.tensor.optimize
+    :members:

pytensor/tensor/__init__.py

@@ -118,6 +118,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
 from pytensor.tensor import linalg
 from pytensor.tensor import special
 from pytensor.tensor import signal
+from pytensor.tensor import optimize
 
 # For backward compatibility
 from pytensor.tensor import nlinalg

pytensor/tensor/optimize.py

@@ -4,17 +4,14 @@
 from typing import cast
 
 import numpy as np
-from scipy.optimize import minimize as scipy_minimize
-from scipy.optimize import minimize_scalar as scipy_minimize_scalar
-from scipy.optimize import root as scipy_root
-from scipy.optimize import root_scalar as scipy_root_scalar
 
 import pytensor.scalar as ps
-from pytensor import Variable, function, graph_replace
+from pytensor.compile.function import function
 from pytensor.gradient import grad, hessian, jacobian
 from pytensor.graph import Apply, Constant, FunctionGraph
 from pytensor.graph.basic import ancestors, truncated_graph_inputs
 from pytensor.graph.op import ComputeMapType, HasInnerGraph, Op, StorageMapType
+from pytensor.graph.replace import graph_replace
 from pytensor.tensor.basic import (
     atleast_2d,
     concatenate,
@@ -24,7 +21,7 @@
 )
 from pytensor.tensor.math import dot
 from pytensor.tensor.slinalg import solve
-from pytensor.tensor.variable import TensorVariable
+from pytensor.tensor.variable import TensorVariable, Variable
 
 
 _log = logging.getLogger(__name__)
@@ -378,6 +375,8 @@ def __init__(
         self._fn_wrapped = None
 
     def perform(self, node, inputs, outputs):
+        from scipy.optimize import minimize_scalar as scipy_minimize_scalar
+
         f = self.fn_wrapped
         f.clear_cache()
 
@@ -426,6 +425,27 @@ def minimize_scalar(
 ):
     """
     Minimize a scalar objective function using scipy.optimize.minimize_scalar.
+
+    Parameters
+    ----------
+    objective : TensorVariable
+        The objective function to minimize. This should be a PyTensor variable representing a scalar value.
+    x : TensorVariable
+        The variable with respect to which the objective function is minimized. It must be a scalar and an
+        input to the computational graph of `objective`.
+    method : str, optional
+        The optimization method to use. Default is "brent". See `scipy.optimize.minimize_scalar` for other options.
+    optimizer_kwargs : dict, optional
+        Additional keyword arguments to pass to `scipy.optimize.minimize_scalar`.
+
+    Returns
+    -------
+    solution: TensorVariable
+        Value of `x` that minimizes `objective(x, *args)`. If the success flag is False, this will be the
+        final state returned by the minimization routine, not necessarily a minimum.
+    success : TensorVariable
+        Symbolic boolean flag indicating whether the minimization routine reported convergence to a minimum
+        value, based on the requested convergence criteria.
     """
 
     args = _find_optimization_parameters(objective, x)
@@ -438,7 +458,11 @@ def minimize_scalar(
         optimizer_kwargs=optimizer_kwargs,
     )
 
-    return minimize_scalar_op(x, *args)
+    solution, success = cast(
+        tuple[TensorVariable, TensorVariable], minimize_scalar_op(x, *args)
+    )
+
+    return solution, success
 
 
 class MinimizeOp(ScipyWrapperOp):
@@ -488,6 +512,8 @@ def __init__(
         self._fn_wrapped = None
 
     def perform(self, node, inputs, outputs):
+        from scipy.optimize import minimize as scipy_minimize
+
         f = self.fn_wrapped
         x0, *args = inputs
 
@@ -538,7 +564,7 @@ def minimize(
     jac: bool = True,
     hess: bool = False,
     optimizer_kwargs: dict | None = None,
-):
+) -> tuple[TensorVariable, TensorVariable]:
     """
     Minimize a scalar objective function using scipy.optimize.minimize.
 
@@ -563,9 +589,13 @@ def minimize(
 
     Returns
     -------
-    TensorVariable
-        The optimized value of x that minimizes the objective function.
+    solution: TensorVariable
+        The optimized value of the vector of inputs `x` that minimizes `objective(x, *args)`. If the success flag
+        is False, this will be the final state of the minimization routine, but not necessarily a minimum.
 
+    success: TensorVariable
+        Symbolic boolean flag indicating whether the minimization routine reported convergence to a minimum
+        value, based on the requested convergence criteria.
     """
     args = _find_optimization_parameters(objective, x)
 
@@ -579,7 +609,11 @@ def minimize(
         optimizer_kwargs=optimizer_kwargs,
     )
 
-    return minimize_op(x, *args)
+    solution, success = cast(
+        tuple[TensorVariable, TensorVariable], minimize_op(x, *args)
+    )
+
+    return solution, success
 
 
 class RootScalarOp(ScipyScalarWrapperOp):
@@ -634,6 +668,8 @@ def __init__(
         self._fn_wrapped = None
 
     def perform(self, node, inputs, outputs):
+        from scipy.optimize import root_scalar as scipy_root_scalar
+
         f = self.fn_wrapped
         f.clear_cache()
         # f.copy_x = True
@@ -676,19 +712,48 @@ def L_op(self, inputs, outputs, output_grads):
 
 def root_scalar(
     equation: TensorVariable,
-    variables: TensorVariable,
+    variable: TensorVariable,
     method: str = "secant",
     jac: bool = False,
     hess: bool = False,
     optimizer_kwargs: dict | None = None,
-):
+) -> tuple[TensorVariable, TensorVariable]:
     """
     Find roots of a scalar equation using scipy.optimize.root_scalar.
+
+    Parameters
+    ----------
+    equation : TensorVariable
+        The equation for which to find roots. This should be a PyTensor variable representing a single equation in one
+        variable. The function will find `variables` such that `equation(variables, *args) = 0`.
+    variable : TensorVariable
+        The variable with respect to which the equation is solved. It must be a scalar and an input to the
+        computational graph of `equation`.
+    method : str, optional
+        The root-finding method to use. Default is "secant". See `scipy.optimize.root_scalar` for other options.
+    jac : bool, optional
+        Whether to compute and use the first derivative of the equation with respect to `variables`.
+        Default is False. Some methods require this.
+    hess : bool, optional
+        Whether to compute and use the second derivative of the equation with respect to `variables`.
+        Default is False. Some methods require this.
+    optimizer_kwargs : dict, optional
+        Additional keyword arguments to pass to `scipy.optimize.root_scalar`.
+
+    Returns
+    -------
+    solution: TensorVariable
+        The final state of the root-finding routine. When `success` is True, this is the value of `variables` that
+        causes `equation` to evaluate to zero. Otherwise it is the final state returned by the root-finding
+        routine, but not necessarily a root.
+
+    success: TensorVariable
+        Boolean indicating whether the root-finding was successful. If True, the solution is a root of the equation
     """
-    args = _find_optimization_parameters(equation, variables)
+    args = _find_optimization_parameters(equation, variable)
 
     root_scalar_op = RootScalarOp(
-        variables,
+        variable,
         *args,
         equation=equation,
         method=method,
@@ -697,7 +762,11 @@ def root_scalar(
         optimizer_kwargs=optimizer_kwargs,
     )
 
-    return root_scalar_op(variables, *args)
+    solution, success = cast(
+        tuple[TensorVariable, TensorVariable], root_scalar_op(variable, *args)
+    )
+
+    return solution, success
 
 
 class RootOp(ScipyWrapperOp):
@@ -761,6 +830,8 @@ def build_fn(self):
         self._fn_wrapped = LRUCache1(fn)
 
     def perform(self, node, inputs, outputs):
+        from scipy.optimize import root as scipy_root
+
         f = self.fn_wrapped
         f.clear_cache()
         f.copy_x = True
@@ -815,8 +886,36 @@ def root(
     method: str = "hybr",
     jac: bool = True,
     optimizer_kwargs: dict | None = None,
-):
-    """Find roots of a system of equations using scipy.optimize.root."""
+) -> tuple[TensorVariable, TensorVariable]:
+    """
+    Find roots of a system of equations using scipy.optimize.root.
+
+    Parameters
+    ----------
+    equations : TensorVariable
+        The system of equations for which to find roots. This should be a PyTensor variable representing a
+        vector (or scalar) value. The function will find `variables` such that `equations(variables, *args) = 0`.
+    variables : TensorVariable
+        The variable(s) with respect to which the system of equations is solved. It must be an input to the
+        computational graph of `equations` and have the same number of dimensions as `equations`.
+    method : str, optional
+        The root-finding method to use. Default is "hybr". See `scipy.optimize.root` for other options.
+    jac : bool, optional
+        Whether to compute and use the Jacobian of the `equations` with respect to `variables`.
+        Default is True. Most methods require this.
+    optimizer_kwargs : dict, optional
+        Additional keyword arguments to pass to `scipy.optimize.root`.
+
+    Returns
+    -------
+    solution: TensorVariable
+        The final state of the root-finding routine. When `success` is True, this is the value of `variables` that
+        causes all `equations` to evaluate to zero. Otherwise it is the final state returned by the root-finding
+        routine, but not necessarily a root.
+
+    success: TensorVariable
+        Boolean indicating whether the root-finding was successful. If True, the solution is a root of the equation
+    """
 
     args = _find_optimization_parameters(equations, variables)
 
@@ -829,7 +928,11 @@ def root(
         optimizer_kwargs=optimizer_kwargs,
     )
 
-    return root_op(variables, *args)
+    solution, success = cast(
+        tuple[TensorVariable, TensorVariable], root_op(variables, *args)
+    )
+
+    return solution, success
 
 
 __all__ = ["minimize_scalar", "minimize", "root_scalar", "root"]

scripts/generate_gallery.py

@@ -58,6 +58,7 @@
     "introduction": "Introduction",
     "rewrites": "Graph Rewriting",
     "scan": "Looping in Pytensor",
+    "optimize": "Optimization in Pytensor",
 }