Change TransformState to NamedTuple (#106)

* Change TransformState to NamedTuple

* Change class type docstring from Args to Attributes

* Update inplace gotcha

* Update docstring

Author: SamDuffield

docs/getting_started.md

@@ -64,7 +64,7 @@ Here:
 - `build` is a function that loads `config_args` into the `init` and `update` functions
  and stores them within the `transform` instance. The `init` and `update` 
  functions then conform to a preset signature allowing for easy switching between algorithms.
-- `state` is a [`dataclass`](https://docs.python.org/3/library/dataclasses.html)
+- `state` is a [`NamedTuple`](https://docs.python.org/3/library/typing.html#typing.NamedTuple)
     encoding the state of the algorithm, including `params` and `aux` attributes.
 - `init` constructs the iteration-varying `state` based on the model parameters `params`.
 - `update` updates the `state` based on a new `batch` of data.

docs/gotchas.md

@@ -70,6 +70,29 @@ state2 = transform.update(state, batch, inplace=True)
 # state is updated and state2 is a pointer to state
 ```
 
+When adding a new algorithm, in-place support can be achieved by modifying `TensorTree`s
+via the [`flexi_tree_map`](https://normal-computing.github.io/posteriors/api/tree_utils/#posteriors.tree_utils.flexi_tree_map) function:
+
+```python
+from posteriors.tree_utils import flexi_tree_map
+
+new_state = flexi_tree_map(lambda x: x + 1, state, inplace=True)
+```
+
+As `posteriors` transform states are immutable `NamedTuple`s, in-place modification of
+`TensorTree` leaves can be achieved by modifying the data of the tensor directly with [`tree_insert_`](https://normal-computing.github.io/posteriors/api/tree_utils/#posteriors.tree_utils.tree_insert_):
+
+```python
+from posteriors.tree_utils import tree_insert_
+
+tree_insert_(state.log_posterior, log_post.detach())
+```
+
+However, the `aux` component of the `TransformState` is not guaranteed to be a `TensorTree`,
+and so in-place modification of `aux` is not supported. Using `state._replace(aux=aux)`
+will return a state with all `TensorTree` pointing to the same memory as input `state`,
+but with a new `aux` component (`aux` is not modified in the input `state` object).
+
 
 ## `torch.tensor` with autograd
 

docs/tutorials/lightning_autoencoder.md

@@ -16,7 +16,6 @@ from torchvision.datasets import MNIST
 from torchvision.transforms import ToTensor
 import lightning as L
 import torchopt
-from dataclasses import asdict
 
 import posteriors
 
@@ -100,7 +99,7 @@ class LitAutoEncoderUQ(L.LightningModule):
         # it is independent of forward
         self.state = self.transform.update(self.state, batch, inplace=True)
         # Logging to TensorBoard (if installed) by default
-        for k, v in asdict(self.state).items():
+        for k, v in self.state._asdict().items():
             if isinstance(v, float) or (isinstance(v, torch.Tensor) and v.numel() == 1):
                 self.log(k, v)
 

examples/lightning_autoencoder.py

@@ -5,7 +5,6 @@
 from torchvision.transforms import ToTensor
 import lightning as L
 import torchopt
-from dataclasses import asdict
 
 import posteriors
 
@@ -54,7 +53,7 @@ def training_step(self, batch, batch_idx):
         # it is independent of forward
         self.state = self.transform.update(self.state, batch, inplace=True)
         # Logging to TensorBoard (if installed) by default
-        for k, v in asdict(self.state).items():
+        for k, v in self.state._asdict().items():
             if isinstance(v, float) or (isinstance(v, torch.Tensor) and v.numel() == 1):
                 self.log(k, v)
 

posteriors/ekf/dense_fisher.py

@@ -1,13 +1,12 @@
-from typing import Any
+from typing import Any, NamedTuple
 from functools import partial
 import torch
 from torch.func import grad_and_value
-from dataclasses import dataclass
 from optree.integration.torch import tree_ravel
 
-from posteriors.tree_utils import tree_size
+from posteriors.tree_utils import tree_size, tree_insert_
 
-from posteriors.types import TensorTree, Transform, LogProbFn, TransformState
+from posteriors.types import TensorTree, Transform, LogProbFn
 from posteriors.utils import (
     per_samplify,
     empirical_fisher,
@@ -67,11 +66,10 @@ def build(
     return Transform(init_fn, update_fn)
 
 
-@dataclass
-class EKFDenseState(TransformState):
+class EKFDenseState(NamedTuple):
     """State encoding a Normal distribution over parameters.
 
-    Args:
+    Attributes:
         params: Mean of the Normal distribution.
         cov: Covariance matrix of the
             Normal distribution.
@@ -81,7 +79,7 @@ class EKFDenseState(TransformState):
 
     params: TensorTree
     cov: torch.Tensor
-    log_likelihood: float = 0
+    log_likelihood: torch.Tensor = torch.tensor([])
     aux: Any = None
 
 
@@ -170,11 +168,11 @@ def log_likelihood_reduced(params, batch):
         update_mean = mu_unravel_f(update_mean)
 
     if inplace:
-        state.params = update_mean
-        state.cov = update_cov
-        state.log_likelihood = log_liks.mean().detach()
-        state.aux = aux
-        return state
+        tree_insert_(state.params, update_mean)
+        tree_insert_(state.cov, update_cov)
+        tree_insert_(state.log_likelihood, log_liks.mean().detach())
+        return state._replace(aux=aux)
+
     return EKFDenseState(update_mean, update_cov, log_liks.mean().detach(), aux)
 
 

posteriors/ekf/diag_fisher.py

@@ -1,12 +1,11 @@
-from typing import Any
+from typing import Any, NamedTuple
 from functools import partial
 import torch
 from torch.func import jacrev
 from optree import tree_map
-from dataclasses import dataclass
 
-from posteriors.types import TensorTree, Transform, LogProbFn, TransformState
-from posteriors.tree_utils import flexi_tree_map
+from posteriors.types import TensorTree, Transform, LogProbFn
+from posteriors.tree_utils import flexi_tree_map, tree_insert_
 from posteriors.utils import (
     diag_normal_sample,
     per_samplify,
@@ -68,11 +67,10 @@ def build(
     return Transform(init_fn, update_fn)
 
 
-@dataclass
-class EKFDiagState(TransformState):
+class EKFDiagState(NamedTuple):
     """State encoding a diagonal Normal distribution over parameters.
 
-    Args:
+    Attributes:
         params: Mean of the Normal distribution.
         sd_diag: Square-root diagonal of the covariance matrix of the
             Normal distribution.
@@ -82,7 +80,7 @@ class EKFDiagState(TransformState):
 
     params: TensorTree
     sd_diag: TensorTree
-    log_likelihood: float = 0
+    log_likelihood: torch.Tensor = torch.tensor([])
     aux: Any = None
 
 
@@ -176,9 +174,9 @@ def update(
     )
 
     if inplace:
-        state.log_likelihood = log_liks.mean().detach()
-        state.aux = aux
-        return state
+        tree_insert_(state.log_likelihood, log_liks.mean().detach())
+        return state._replace(aux=aux)
+
     return EKFDiagState(update_mean, update_sd_diag, log_liks.mean().detach(), aux)
 
 

posteriors/laplace/dense_fisher.py

@@ -1,11 +1,10 @@
-from typing import Any
-from dataclasses import dataclass
+from typing import Any, NamedTuple
 from functools import partial
 import torch
 from optree import tree_map
 from optree.integration.torch import tree_ravel
 
-from posteriors.types import TensorTree, Transform, LogProbFn, TransformState
+from posteriors.types import TensorTree, Transform, LogProbFn
 from posteriors.tree_utils import tree_size
 from posteriors.utils import (
     per_samplify,
@@ -55,12 +54,11 @@ def build(
     return Transform(init_fn, update_fn)
 
 
-@dataclass
-class DenseLaplaceState(TransformState):
+class DenseLaplaceState(NamedTuple):
     """State encoding a Normal distribution over parameters,
     with a dense precision matrix
 
-    Args:
+    Attributes:
         params: Mean of the Normal distribution.
         prec: Precision matrix of the Normal distribution.
         aux: Auxiliary information from the log_posterior call.
@@ -130,9 +128,8 @@ def update(
         )(state.params)
 
     if inplace:
-        state.prec += fisher
-        state.aux = aux
-        return state
+        state.prec.data += fisher
+        return state._replace(aux=aux)
     else:
         return DenseLaplaceState(state.params, state.prec + fisher, aux)
 

posteriors/laplace/dense_ggn.py

@@ -1,16 +1,14 @@
 from functools import partial
-from typing import Any
+from typing import Any, NamedTuple
 import torch
 from optree import tree_map
-from dataclasses import dataclass
 from optree.integration.torch import tree_ravel
 
 from posteriors.types import (
     TensorTree,
     Transform,
     ForwardFn,
     OuterLogProbFn,
-    TransformState,
 )
 from posteriors.utils import (
     tree_size,
@@ -67,12 +65,11 @@ def build(
     return Transform(init_fn, update_fn)
 
 
-@dataclass
-class DenseLaplaceState(TransformState):
+class DenseLaplaceState(NamedTuple):
     """State encoding a Normal distribution over parameters,
     with a dense precision matrix
 
-    Args:
+    Attributes:
         params: Mean of the Normal distribution.
         prec: Precision matrix of the Normal distribution.
         aux: Auxiliary information from the log_posterior call.
@@ -145,9 +142,8 @@ def outer_loss(z, batch):
         )(state.params)
 
     if inplace:
-        state.prec += ggn_batch
-        state.aux = aux
-        return state
+        state.prec.data += ggn_batch
+        return state._replace(aux=aux)
     else:
         return DenseLaplaceState(state.params, state.prec + ggn_batch, aux)
 

posteriors/laplace/diag_fisher.py

@@ -1,11 +1,10 @@
 from functools import partial
-from typing import Any
+from typing import Any, NamedTuple
 import torch
 from torch.func import jacrev
 from optree import tree_map
-from dataclasses import dataclass
 
-from posteriors.types import TensorTree, Transform, LogProbFn, TransformState
+from posteriors.types import TensorTree, Transform, LogProbFn
 from posteriors.tree_utils import flexi_tree_map
 from posteriors.utils import (
     diag_normal_sample,
@@ -54,11 +53,10 @@ def build(
     return Transform(init_fn, update_fn)
 
 
-@dataclass
-class DiagLaplaceState(TransformState):
+class DiagLaplaceState(NamedTuple):
     """State encoding a diagonal Normal distribution over parameters.
 
-    Args:
+    Attributes:
         params: Mean of the Normal distribution.
         prec_diag: Diagonal of the precision matrix of the Normal distribution.
         aux: Auxiliary information from the log_posterior call.
@@ -134,8 +132,7 @@ def update_func(x, y):
     )
 
     if inplace:
-        state.aux = aux
-        return state
+        return state._replace(aux=aux)
     return DiagLaplaceState(state.params, prec_diag, aux)
 
 

posteriors/laplace/diag_ggn.py

@@ -1,15 +1,13 @@
 from functools import partial
-from typing import Any
+from typing import Any, NamedTuple
 import torch
 from optree import tree_map
-from dataclasses import dataclass
 
 from posteriors.types import (
     TensorTree,
     Transform,
     ForwardFn,
     OuterLogProbFn,
-    TransformState,
 )
 from posteriors.tree_utils import flexi_tree_map
 from posteriors.utils import (
@@ -66,11 +64,10 @@ def build(
     return Transform(init_fn, update_fn)
 
 
-@dataclass
-class DiagLaplaceState(TransformState):
+class DiagLaplaceState(NamedTuple):
     """State encoding a diagonal Normal distribution over parameters.
 
-    Args:
+    Attributes:
         params: Mean of the Normal distribution.
         prec_diag: Diagonal of the precision matrix of the Normal distribution.
         aux: Auxiliary information from the log_posterior call.
@@ -149,8 +146,7 @@ def update_func(x, y):
     )
 
     if inplace:
-        state.aux = aux
-        return state
+        return state._replace(aux=aux)
     return DiagLaplaceState(state.params, prec_diag, aux)