Changes to support float32 inputs

Author: jessegrabowski

pytensor/tensor/optimize.py

@@ -52,6 +52,9 @@ def __init__(self, fn, copy_x: bool = False):
         self.last_result = None
         self.copy_x = copy_x
 
+        # Scipy does not respect dtypes *at all*, so we have to force it ourselves.
+        self.dtype = fn.maker.fgraph.inputs[0].type.dtype
+
         self.cache_hits = 0
         self.cache_misses = 0
 
@@ -67,9 +70,7 @@ def __call__(self, x, *args):
         If the input `x` is the same as the last input, return the cached result. Otherwise update the cache with the
         new input and result.
         """
-        # scipy.optimize.scalar_minimize and scalar_root don't take initial values as an argument, so we can't control
-        # the first input to the inner function. Of course, they use a scalar, but we need a 0d numpy array.
-        x = np.asarray(x)
+        x = x.astype(self.dtype)
 
         if self.last_result is None or not (x == self.last_x).all():
             self.cache_misses += 1
@@ -160,6 +161,7 @@ def _get_parameter_grads_from_vector(
         )
 
         grad_wrt_args.append(dot(output_grad, arg_grad))
+
         cursor += arg_size
 
     return grad_wrt_args
@@ -175,17 +177,11 @@ def build_fn(self):
         """
         outputs = self.inner_outputs
         self._fn = fn = function(self.inner_inputs, outputs, trust_input=True)
+
         # Do this reassignment to see the compiled graph in the dprint
         # self.fgraph = fn.maker.fgraph
 
-        if self.inner_inputs[0].type.shape == ():
-
-            def fn_wrapper(x, *args):
-                return fn(x.squeeze(), *args)
-
-            self._fn_wrapped = LRUCache1(fn_wrapper)
-        else:
-            self._fn_wrapped = LRUCache1(fn)
+        self._fn_wrapped = LRUCache1(fn)
 
     @property
     def fn(self):
@@ -771,7 +767,9 @@ def perform(self, node, inputs, outputs):
             **self.optimizer_kwargs,
         )
 
-        outputs[0][0] = res.x.reshape(variables.shape)
+        # There's a reshape here to cover the case where variables is a scalar. Scipy will still return a
+        # (1, 1) matrix in in this case, which causes errors downstream (since pytensor expects a scalar).
+        outputs[0][0] = res.x.reshape(variables.shape).astype(variables.dtype)
         outputs[1][0] = np.bool_(res.success)
 
     def L_op(
@@ -807,12 +805,20 @@ def root(
     variables: TensorVariable,
     method: str = "hybr",
     jac: bool = True,
+    optimizer_kwargs: dict | None = None,
 ):
     """Find roots of a system of equations using scipy.optimize.root."""
 
     args = _find_optimization_parameters(equations, variables)
 
-    root_op = RootOp(variables, *args, equations=equations, method=method, jac=jac)
+    root_op = RootOp(
+        variables,
+        *args,
+        equations=equations,
+        method=method,
+        jac=jac,
+        optimizer_kwargs=optimizer_kwargs,
+    )
 
     return root_op(variables, *args)
 

tests/tensor/test_optimize.py

@@ -59,7 +59,12 @@ def test_simple_minimize():
     minimized_x_val, success_val = f(a_val, c_val, 0.0)
 
     assert success_val
-    assert minimized_x_val == (2 * a_val * c_val)
+    np.testing.assert_allclose(
+        minimized_x_val,
+        2 * a_val * c_val,
+        atol=1e-8 if config.floatX == "float64" else 1e-6,
+        rtol=1e-8 if config.floatX == "float64" else 1e-6,
+    )
 
     def f(x, a, b):
         objective = (x - a * b) ** 2
@@ -82,7 +87,7 @@ def test_minimize_vector_x(method, jac, hess):
     def rosenbrock_shifted_scaled(x, a, b):
         return (a * (x[1:] - x[:-1] ** 2) ** 2 + (1 - x[:-1]) ** 2).sum() + b
 
-    x = pt.dvector("x")
+    x = pt.tensor("x", shape=(None,))
     a = pt.scalar("a")
     b = pt.scalar("b")
 
@@ -91,23 +96,30 @@ def rosenbrock_shifted_scaled(x, a, b):
         objective, x, method=method, jac=jac, hess=hess, optimizer_kwargs={"tol": 1e-16}
     )
 
-    a_val = 0.5
-    b_val = 1.0
-    x0 = np.zeros(5).astype(floatX)
-    x_star_val = minimized_x.eval({a: a_val, b: b_val, x: x0})
+    fn = pytensor.function([x, a, b], [minimized_x, success])
 
-    assert success.eval({a: a_val, b: b_val, x: x0})
+    a_val = np.array(0.5, dtype=floatX)
+    b_val = np.array(1.0, dtype=floatX)
+    x0 = np.zeros((5,)).astype(floatX)
+    x_star_val, success = fn(x0, a_val, b_val)
+
+    assert success
 
     np.testing.assert_allclose(
-        x_star_val, np.ones_like(x_star_val), atol=1e-6, rtol=1e-6
+        x_star_val,
+        np.ones_like(x_star_val),
+        atol=1e-8 if config.floatX == "float64" else 1e-3,
+        rtol=1e-8 if config.floatX == "float64" else 1e-3,
     )
 
+    assert x_star_val.dtype == floatX
+
     def f(x, a, b):
         objective = rosenbrock_shifted_scaled(x, a, b)
         out = minimize(objective, x)[0]
         return out
 
-    utt.verify_grad(f, [x0, a_val, b_val], eps=1e-6)
+    utt.verify_grad(f, [x0, a_val, b_val], eps=1e-3 if floatX == "float32" else 1e-6)
 
 
 @pytest.mark.parametrize(
@@ -130,7 +142,12 @@ def fn(x, a):
     solution, success = func(x0, a_val)
 
     assert success
-    np.testing.assert_allclose(solution, -1.02986653, atol=1e-6, rtol=1e-6)
+    np.testing.assert_allclose(
+        solution,
+        -1.02986653,
+        atol=1e-8 if config.floatX == "float64" else 1e-6,
+        rtol=1e-8 if config.floatX == "float64" else 1e-6,
+    )
 
     def root_fn(x, a):
         f = fn(x, a)
@@ -147,15 +164,20 @@ def fn(x, a):
         return x + 2 * a * pt.cos(x)
 
     f = fn(x, a)
-    root_f, success = root(f, x)
+    root_f, success = root(f, x, method="lm", optimizer_kwargs={"tol": 1e-8})
     func = pytensor.function([x, a], [root_f, success])
 
     x0 = 0.0
     a_val = 1.0
     solution, success = func(x0, a_val)
 
     assert success
-    np.testing.assert_allclose(solution, -1.02986653, atol=1e-6, rtol=1e-6)
+    np.testing.assert_allclose(
+        solution,
+        -1.02986653,
+        atol=1e-8 if config.floatX == "float64" else 1e-6,
+        rtol=1e-8 if config.floatX == "float64" else 1e-6,
+    )
 
     def root_fn(x, a):
         f = fn(x, a)
@@ -165,24 +187,27 @@ def root_fn(x, a):
 
 
 def test_root_system_of_equations():
-    x = pt.dvector("x")
-    a = pt.dvector("a")
-    b = pt.dvector("b")
+    x = pt.tensor("x", shape=(None,))
+    a = pt.tensor("a", shape=(None,))
+    b = pt.tensor("b", shape=(None,))
 
     f = pt.stack([a[0] * x[0] * pt.cos(x[1]) - b[0], x[0] * x[1] - a[1] * x[1] - b[1]])
 
-    root_f, success = root(f, x)
+    root_f, success = root(f, x, method="lm", optimizer_kwargs={"tol": 1e-8})
     func = pytensor.function([x, a, b], [root_f, success])
 
-    x0 = np.array([1.0, 1.0])
-    a_val = np.array([1.0, 1.0])
-    b_val = np.array([4.0, 5.0])
+    x0 = np.array([1.0, 1.0], dtype=floatX)
+    a_val = np.array([1.0, 1.0], dtype=floatX)
+    b_val = np.array([4.0, 5.0], dtype=floatX)
     solution, success = func(x0, a_val, b_val)
 
     assert success
 
     np.testing.assert_allclose(
-        solution, np.array([6.50409711, 0.90841421]), atol=1e-6, rtol=1e-6
+        solution,
+        np.array([6.50409711, 0.90841421]),
+        atol=1e-8 if config.floatX == "float64" else 1e-6,
+        rtol=1e-8 if config.floatX == "float64" else 1e-6,
     )
 
     def root_fn(x, a, b):
@@ -191,4 +216,6 @@ def root_fn(x, a, b):
         )
         return root(f, x)[0]
 
-    utt.verify_grad(root_fn, [x0, a_val, b_val], eps=1e-6)
+    utt.verify_grad(
+        root_fn, [x0, a_val, b_val], eps=1e-6 if floatX == "float64" else 1e-3
+    )