Raise pytorch dependency to 2.0

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## 0.2.0 - 2023-11-10
+
+### Changed
+
+- Removed compatibility with pytorch 1.x
+
 ## 0.1.9 - 2023-08-29
 
 ### Added

README.md

@@ -53,9 +53,6 @@ step_size_controller = to.IntegralController(atol=1e-6, rtol=1e-3, term=term)
 solver = to.AutoDiffAdjoint(step_method, step_size_controller)
 jit_solver = torch.compile(solver)
 
-# For pytorch versions < 2.0, use the older TorchScript compiler
-#jit_solver = torch.jit.script(solver)
-
 sol = jit_solver.solve(to.InitialValueProblem(y0=y0, t_eval=t_eval))
 print(sol.stats)
 # => {'n_f_evals': tensor([26, 26]), 'n_steps': tensor([4, 2]),

pyproject.toml

@@ -19,11 +19,8 @@ classifiers = [
 ]
 dependencies = [
   "sympy ~= 1.10",
-  "torch >= 1.11",
+  "torch >= 2.0",
   "torchtyping ~= 0.1.4",
-  # functorch has been integrated into pytorch and this is just a dummy dependency for
-  # compatibility with pytorch <= 1.13
-  "functorch",
 ]
 
 [project.optional-dependencies]

tests/torch_script_test.py

@@ -9,19 +9,12 @@
 torch_version = version.parse(torch.__version__)
 
 
-def compile(module):
-    if torch_version.major < 2:
-        return torch.jit.script(module)
-    else:
-        return torch.compile(module)
-
-
 @pytest.mark.parametrize("step_method", [Dopri5, Heun, Tsit5, Euler])
 def test_can_be_jitted_with_torch_script(step_method):
     _, term, problem = get_problem("sine", [[0.1, 0.15, 1.0], [1.0, 1.9, 2.0]])
     step_size_controller = IntegralController(1e-3, 1e-3, term=term)
     adjoint = AutoDiffAdjoint(step_method(term), step_size_controller)
-    jitted = compile(adjoint)
+    jitted = torch.compile(adjoint)
 
     dt0 = torch.tensor([0.01, 0.01]) if step_method is Euler else None
     solution = adjoint.solve(problem, dt0=dt0)
@@ -33,9 +26,9 @@ def test_can_be_jitted_with_torch_script(step_method):
 
 methods = [Dopri5, Heun, Tsit5]
 v = torch_version
-if (v.major, v.minor) not in [(1, 13), (2, 0)]:
-    # In pytorch 1.13.0 and 2.0, Euler triggers an internal error in the JIT compiler
-    # specifically in this next test, so we just exclude it
+if (v.major, v.minor) not in [(2, 0)]:
+    # In pytorch 2.0, Euler triggers an internal error in the JIT compiler specifically
+    # in this next test, so we just exclude it
     methods.append(Euler)
 
 
@@ -51,7 +44,7 @@ def test_passing_term_dynamically_equals_fixed_term(step_method):
 
     controller_jit = IntegralController(1e-3, 1e-3, term=term)
     adjoint_jit = AutoDiffAdjoint(step_method(term), controller_jit)
-    solution_jit = compile(adjoint_jit).solve(problem, dt0=dt0)
+    solution_jit = torch.compile(adjoint_jit).solve(problem, dt0=dt0)
 
     assert solution.ts == approx(solution_jit.ts)
     assert solution.ys == approx(solution_jit.ys, abs=1e-3, rel=1e-3)

torchode/__init__.py

@@ -1,6 +1,6 @@
 """A parallel ODE solver for PyTorch"""
 
-__version__ = "0.1.9"
+__version__ = "0.2.0"
 
 from .adjoints import AutoDiffAdjoint, BacksolveAdjoint, JointBacksolveAdjoint
 from .interface import register_method, solve_ivp