torchode 0.1.9__tar.gz → 0.2.0__tar.gz

{torchode-0.1.9 → torchode-0.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

{torchode-0.1.9 → torchode-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchode
-Version: 0.1.9
+Version: 0.2.0
 Summary: A parallel ODE solver for PyTorch
 Author-email: Marten Lienen <m.lienen@tum.de>
 Requires-Python: >= 3.8
@@ -12,9 +12,8 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Requires-Dist: sympy ~= 1.10
-Requires-Dist: torch >= 1.11
+Requires-Dist: torch >= 2.0
 Requires-Dist: torchtyping ~= 0.1.4
-Requires-Dist: functorch
 Requires-Dist: pytest-benchmark ; extra == "benchmark"
 Requires-Dist: pytest ; extra == "test"
 Requires-Dist: numpy ~= 1.22 ; extra == "test"
@@ -77,9 +76,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]),

{torchode-0.1.9 → torchode-0.2.0}/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]),

{torchode-0.1.9 → torchode-0.2.0}/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]

{torchode-0.1.9 → torchode-0.2.0}/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

{torchode-0.1.9 → torchode-0.2.0}/torchode/adjoints.py

@@ -1,6 +1,5 @@
 from typing import Any, Dict, Optional
 
-import functorch
 import torch
 import torch.nn as nn
 
@@ -449,22 +448,24 @@ class AugmentedDynamicsTerm(nn.Module):
         self.vmap_args_dims = vmap_args_dims
         self.vmap_randomness = vmap_randomness
 
-        self.func, self.params, self.buffers = functorch.make_functional_with_buffers(
-            self.term.f
-        )
+        f = self.term.f
+        orig_params = dict(f.named_parameters())
+        buffers = dict(f.named_buffers())
+        params_names = orig_params.keys()
 
         def vjp_single_sample(t_i, y_i, adj_y_i, arg_i):
             def wrapper(params, t_, y_):
+                args = (t_, y_)
                 if self.term.with_args:
-                    return self.func(params, self.buffers, t_, y_, arg_i)
-                else:
-                    return self.func(params, self.buffers, t_, y_)
+                    args = args + (arg_i,)
+                params_dict = {name: value for name, value in zip(params_names, params)}
+                return torch.func.functional_call(f, (params_dict, buffers), args)
 
-            dy, vjp = functorch.vjp(wrapper, self.params, t_i, y_i)
+            dy, vjp = torch.func.vjp(wrapper, tuple(orig_params.values()), t_i, y_i)
             vjp_params, vjp_t, vjp_y = vjp(-adj_y_i)
             return dy, vjp_t, vjp_y, vjp_params
 
-        self.vjp_vf = functorch.vmap(
+        self.vjp_vf = torch.func.vmap(
             vjp_single_sample,
             in_dims=(0, 0, 0, self.vmap_args_dims),
             randomness=self.vmap_randomness,