cleanup

Author: poonai

diffrax/__init__.py

@@ -45,7 +45,7 @@
     AbstractLocalInterpolation as AbstractLocalInterpolation,
     FourthOrderPolynomialInterpolation as FourthOrderPolynomialInterpolation,
     LocalLinearInterpolation as LocalLinearInterpolation,
-    RodasInterpolation as RodasInterpolation,  # noqa: E501
+    RodasInterpolation as RodasInterpolation,
     ThirdOrderHermitePolynomialInterpolation as ThirdOrderHermitePolynomialInterpolation,  # noqa: E501
 )
 from ._misc import adjoint_rms_seminorm as adjoint_rms_seminorm

diffrax/_local_interpolation.py

@@ -199,12 +199,17 @@ def evaluate(
             return self.evaluate(t1) - self.evaluate(t0)
 
         t = linear_rescale(self.t0, t0, self.t1)
-        weighted_increment = jax.vmap(
-            lambda coeff, stage_k: (t * jnp.polyval(jnp.flip(coeff), t)) * stage_k
-        )(self.stage_poly_coeffs, self.k)
+
+        def eval_increment():
+            with jax.numpy_dtype_promotion("standard"):
+                weighted_increment = jax.vmap(
+                    lambda coeff, stage_k: (t * jnp.polyval(jnp.flip(coeff), t))
+                    * stage_k
+                )(self.stage_poly_coeffs, self.k)
+                return jnp.sum(weighted_increment, axis=0).astype(self.k.dtype)
 
         y0, unravel = fu.ravel_pytree(self.y0)
-        y1 = y0 + jnp.sum(weighted_increment, axis=0)
+        y1 = y0 + eval_increment()
         return unravel(y1)
 
     @classmethod

diffrax/_solver/rodas5p.py

@@ -1,6 +1,5 @@
 from collections.abc import Callable
 from typing import ClassVar
-from xmlrpc.client import Boolean
 
 import numpy as np
 
@@ -229,7 +228,7 @@ class Rodas5p(AbstractRosenbrock):
         _Rodas5pInterpolation.from_k
     )
 
-    rodas: ClassVar[Boolean] = True
+    rodas: ClassVar[bool] = True
 
     def order(self, terms):
         del terms

diffrax/_solver/rosenbrock.py

@@ -101,7 +101,7 @@ class AbstractRosenbrock(AbstractAdaptiveSolver):
 
     rodas: ClassVar[bool] = False
 
-    linear_solver: lx.AbstractLinearSolver = lx.LU()
+    linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=True)
 
     def init(self, terms, t0, t1, y0, args) -> _SolverState:
         del t0, t1

test/test_integrate.py

@@ -152,7 +152,7 @@ def test_ode_order(solver, dtype):
     A = jr.normal(akey, (10, 10), dtype=dtype) * 0.5
 
     if isinstance(solver, AbstractRosenbrock) and dtype == jnp.complex128:
-        ## complex support is not added to rosenbrock.
+        # complex support is not added to rosenbrock.
         return
 
     if (

test/test_solver.py

@@ -465,7 +465,6 @@ def rober(t, y, args):
             [6.1723488239606716e-01, 6.1535912746388841e-06, 3.8275896401264059e-01],
         ]
     )
-    print(sol.ys)
     assert jnp.allclose(sol.ys, true_ys, rtol=1e-3, atol=1e-8)  # pyright: ignore