Fix optree deprecation warning (#144)

Author: SamDuffield

.pre-commit-config.yaml

@@ -1,7 +1,10 @@
 repos:
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.5
-    hooks:
-    - id: ruff
-      args: ["--fix"]
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.13.1
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+      args: [ --fix ]
+    # Run the formatter.
     - id: ruff-format

posteriors/ekf/dense_fisher.py

@@ -2,7 +2,7 @@
 from functools import partial
 import torch
 from torch.func import grad_and_value
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 from tensordict import TensorClass
 
 from posteriors.tree_utils import tree_size, tree_insert_

posteriors/laplace/dense_fisher.py

@@ -2,7 +2,7 @@
 from functools import partial
 import torch
 from optree import tree_map
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 from tensordict import TensorClass
 from posteriors.types import TensorTree, Transform, LogProbFn
 from posteriors.tree_utils import tree_size, tree_insert_

posteriors/laplace/dense_ggn.py

@@ -2,7 +2,7 @@
 from typing import Any
 import torch
 from optree import tree_map
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 from tensordict import TensorClass
 
 from posteriors.types import (

posteriors/laplace/dense_hessian.py

@@ -2,7 +2,7 @@
 from functools import partial
 import torch
 from optree import tree_map
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 from tensordict import TensorClass
 
 from posteriors.types import TensorTree, Transform, LogProbFn

posteriors/sgmcmc/sgnht.py

@@ -3,7 +3,7 @@
 import torch
 from torch.func import grad_and_value
 from optree import tree_map
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 from tensordict import TensorClass
 from posteriors.types import TensorTree, Transform, LogProbFn, Schedule
 from posteriors.tree_utils import flexi_tree_map, tree_insert_

posteriors/utils.py

@@ -6,7 +6,7 @@
 from torch.func import grad, jvp, vjp, functional_call, jacrev, jacfwd
 from torch.distributions import Normal
 from optree import tree_map, tree_reduce, tree_flatten, tree_leaves
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 
 from posteriors.types import TensorTree, ForwardFn, Tensor
 from posteriors.tree_utils import tree_size

posteriors/vi/dense.py

@@ -3,7 +3,7 @@
 import torch
 from torch.func import grad_and_value, vmap
 from optree import tree_map
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 import torchopt
 from tensordict import TensorClass
 

pyproject.toml

@@ -14,7 +14,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Mathematics",
     "License :: OSI Approved :: Apache Software License",
 ]
-dependencies = ["torch>=2.0.0", "torchopt>=0.7.3", "optree>=0.10.0", "tensordict>=0.7.0"]
+dependencies = ["torch>=2.0.0", "torchopt>=0.7.3", "optree>=0.17.0", "tensordict>=0.7.0"]
 
 [project.optional-dependencies]
 test = ["pre-commit", "pytest-cov", "pytest-xdist", "ruff"]

tests/ekf/test_diag_fisher.py

@@ -1,6 +1,6 @@
 from functools import partial
 import torch
-from optree.integration.torch import tree_ravel
+from optree.integrations.torch import tree_ravel
 from posteriors import ekf
 from tests.scenarios import get_multivariate_normal_log_prob
 from tests.utils import verify_inplace_update