ml4gw 0.7.7__py3-none-any.whl → 0.7.8__py3-none-any.whl

ml4gw/augmentations.py

@@ -1,3 +1,8 @@
+"""
+This module contains transformations that may be useful
+for augmenting timeseries data during training
+"""
+
 import torch
 from jaxtyping import Float
 from torch import Tensor

ml4gw/dataloading/__init__.py

@@ -1,3 +1,8 @@
+"""
+This module contains tools for efficient in-memory and
+out-of-memory dataloading.
+"""
+
 from .chunked_dataset import ChunkedTimeSeriesDataset
 from .hdf5_dataset import Hdf5TimeSeriesDataset
 from .in_memory_dataset import InMemoryDataset

ml4gw/dataloading/chunked_dataset.py

@@ -94,10 +94,8 @@ class ChunkedTimeSeriesDataset(torch.utils.data.IterableDataset):
             # flatten it to make it easier to slice
             if chunk_size < self.kernel_size:
                 raise ValueError(
-                    (
-                        "Can't sample kernels of size {} from chunk "
-                        "with size {}"
-                    ).format(self.kernel_size, chunk_size)
+                    f"Can't sample kernels of size {self.kernel_size} from "
+                    f"chunk with size {chunk_size}"
                 )
             chunk = chunk.reshape(-1)
 

ml4gw/dataloading/hdf5_dataset.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Optional, Sequence, Union
+from collections.abc import Sequence
 
 import h5py
 import numpy as np
@@ -63,13 +63,13 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
         kernel_size: int,
         batch_size: int,
         batches_per_epoch: int,
-        coincident: Union[bool, str],
-        num_files_per_batch: Optional[int] = None,
+        coincident: bool | str,
+        num_files_per_batch: int | None = None,
     ) -> None:
         if not isinstance(coincident, bool) and coincident != "files":
             raise ValueError(
                 "coincident must be either a boolean or 'files', "
-                "got unrecognized value {}".format(coincident)
+                f"got unrecognized value {coincident}"
             )
 
         self.fnames = np.array(fnames)
@@ -94,13 +94,11 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
                 dset = f[channels[0]]
                 if dset.chunks is None:
                     warnings.warn(
-                        "File {} contains datasets that were generated "
+                        f"File {fname} contains datasets that were generated "
                         "without using chunked storage. This can have "
                         "severe performance impacts at data loading time. "
                         "If you need faster loading, try re-generating "
-                        "your dataset with chunked storage turned on.".format(
-                            fname
-                        ),
+                        "your dataset with chunked storage turned on.",
                         category=ContiguousHdf5Warning,
                         stacklevel=2,
                     )
@@ -153,7 +151,9 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
         unique_fnames, inv, counts = np.unique(
             fnames, return_inverse=True, return_counts=True
         )
-        for i, (fname, count) in enumerate(zip(unique_fnames, counts)):
+        for i, (fname, count) in enumerate(
+            zip(unique_fnames, counts, strict=True)
+        ):
             size = self.sizes[fname]
             max_idx = size - self.kernel_size
 
@@ -185,7 +185,9 @@ class Hdf5TimeSeriesDataset(torch.utils.data.IterableDataset):
             # open the file and sample a different set of
             # kernels for each batch element it occupies
             with h5py.File(fname, "r") as f:
-                for b, c, i in zip(batch_indices, channel_indices, idx):
+                for b, c, i in zip(
+                    batch_indices, channel_indices, idx, strict=True
+                ):
                     x[b, c] = f[self.channels[c]][i : i + self.kernel_size]
         return torch.Tensor(x)
 

ml4gw/dataloading/in_memory_dataset.py

@@ -1,5 +1,4 @@
 import itertools
-from typing import Optional, Tuple, Union
 
 import torch
 from jaxtyping import Float
@@ -79,10 +78,10 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
         self,
         X: Float[Tensor, "channels time"],
         kernel_size: int,
-        y: Optional[Float[Tensor, " time"]] = None,
+        y: Float[Tensor, " time"] | None = None,
         batch_size: int = 32,
         stride: int = 1,
-        batches_per_epoch: Optional[int] = None,
+        batches_per_epoch: int | None = None,
         coincident: bool = True,
         shuffle: bool = True,
         device: str = "cpu",
@@ -122,10 +121,9 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
             batch_size * batches_per_epoch
         ):
             raise ValueError(
-                "Number of kernels {} in timeseries insufficient "
-                "to generate {} batches of size {}".format(
-                    self.num_kernels, batch_size, batches_per_epoch
-                )
+                f"Number of kernels {self.num_kernels} in timeseries "
+                f"insufficient to generate {batch_size} batches of size "
+                f"{batches_per_epoch}"
             )
 
         self.batch_size = batch_size
@@ -191,7 +189,9 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
             # indices we'll need rather than having to generate
             # everything.
             idx = [range(self.num_kernels) for _ in range(len(self.X))]
-            idx = zip(range(num_kernels), itertools.product(*idx))
+            idx = zip(
+                range(num_kernels), itertools.product(*idx), strict=False
+            )
             idx = torch.stack([torch.Tensor(i[1]) for i in idx])
             idx = idx.type(torch.int64).to(device)
         elif self.shuffle:
@@ -208,10 +208,10 @@ class InMemoryDataset(torch.utils.data.IterableDataset):
 
     def __iter__(
         self,
-    ) -> Union[
-        Float[Tensor, "batch channel time"],
-        Tuple[Float[Tensor, "batch channel time"], Float[Tensor, " batch"]],
-    ]:
+    ) -> (
+        Float[Tensor, "batch channel time"]
+        | tuple[Float[Tensor, "batch channel time"], Float[Tensor, " batch"]]
+    ):
         indices = self.init_indices()
         for i in range(len(self)):
             # slice the array of _indices_ we'll be using to

ml4gw/distributions.py

@@ -6,7 +6,7 @@ from the corresponding distribution.
 """
 
 import math
-from typing import Callable, Optional
+from collections.abc import Callable
 
 import torch
 import torch.distributions as dist
@@ -104,7 +104,7 @@ class LogNormal(dist.LogNormal):
         self,
         mean: float,
         std: float,
-        low: Optional[float] = None,
+        low: float | None = None,
         validate_args=None,
     ):
         self.low = low

ml4gw/gw.py

@@ -1,6 +1,7 @@
 """
-Tools for manipulating raw gravitational waveforms
-and projecting them onto interferometer responses.
+Tools for manipulating raw gravitational waveforms,
+projecting them onto interferometer responses, and
+calculating SNRs.
 Much of the projection code is an extension of the
 implementation made available in
 `bilby <https://arxiv.org/abs/1811.02042>`_.
@@ -8,8 +9,6 @@ Specifically code from
 `this module <https://github.com/lscsoft/bilby/blob/master/bilby/gw/detector/interferometer.py>`_.
 """  # noqa E501
 
-from typing import List, Tuple, Union
-
 import torch
 from jaxtyping import Float
 from torch import Tensor
@@ -58,7 +57,7 @@ def compute_antenna_responses(
     psi: BatchTensor,
     phi: BatchTensor,
     detector_tensors: NetworkDetectorTensors,
-    modes: List[str],
+    modes: list[str],
 ) -> Float[Tensor, "batch polarizations num_ifos"]:
     """
     Compute the antenna pattern factors of a batch of
@@ -257,7 +256,7 @@ def compute_observed_strain(
 
 def get_ifo_geometry(
     *ifos: str,
-) -> Tuple[NetworkDetectorTensors, NetworkVertices]:
+) -> tuple[NetworkDetectorTensors, NetworkVertices]:
     """
     For a given list of interferometer names, retrieve and
     concatenate the associated detector tensors and vertices
@@ -286,8 +285,8 @@ def compute_ifo_snr(
     responses: WaveformTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
-    lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
+    highpass: float | Float[Tensor, " frequency"] | None = None,
+    lowpass: float | Float[Tensor, " frequency"] | None = None,
 ) -> Float[Tensor, "batch num_ifos"]:
     """Compute the SNRs of a batch of interferometer responses
 
@@ -367,10 +366,9 @@ def compute_ifo_snr(
             highpass = freqs >= highpass
         elif len(highpass) != integrand.shape[-1]:
             raise ValueError(
-                "Can't apply highpass filter mask with {} frequency bins"
-                "to signal fft with {} frequency bins".format(
-                    len(highpass), integrand.shape[-1]
-                )
+                f"Can't apply highpass filter mask with {len(highpass)} "
+                f"frequency bins to signal fft with {integrand.shape[-1]} "
+                "frequency bins"
             )
         integrand *= highpass.to(integrand.device)
     if lowpass is not None:
@@ -379,10 +377,9 @@ def compute_ifo_snr(
             lowpass = freqs < lowpass
         elif len(lowpass) != integrand.shape[-1]:
             raise ValueError(
-                "Can't apply lowpass filter mask with {} frequency bins"
-                "to signal fft with {} frequency bins".format(
-                    len(lowpass), integrand.shape[-1]
-                )
+                f"Can't apply lowpass filter mask with {len(lowpass)} "
+                f"frequency bins to signal fft with {integrand.shape[-1]} "
+                "frequency bins"
             )
         integrand *= lowpass.to(integrand.device)
 
@@ -410,8 +407,8 @@ def compute_network_snr(
     responses: WaveformTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
-    lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
+    highpass: float | Float[Tensor, " frequency"] | None = None,
+    lowpass: float | Float[Tensor, " frequency"] | None = None,
 ) -> BatchTensor:
     """
     Compute the total SNR from a gravitational waveform
@@ -467,11 +464,11 @@ def compute_network_snr(
 
 def reweight_snrs(
     responses: WaveformTensor,
-    target_snrs: Union[float, BatchTensor],
+    target_snrs: float | BatchTensor,
     psd: PSDTensor,
     sample_rate: float,
-    highpass: Union[float, Float[Tensor, " frequency"], None] = None,
-    lowpass: Union[float, Float[Tensor, " frequency"], None] = None,
+    highpass: float | Float[Tensor, " frequency"] | None = None,
+    lowpass: float | Float[Tensor, " frequency"] | None = None,
 ) -> WaveformTensor:
     """Scale interferometer responses such that they have a desired SNR
 

ml4gw/nn/__init__.py

@@ -0,0 +1,6 @@
+"""
+This module contains neural network architectures and
+architecture components. These can be a good place
+to get started, rather than defining your own
+architecture from the start.
+"""

ml4gw/nn/autoencoder/base.py

@@ -1,5 +1,4 @@
 from collections.abc import Sequence
-from typing import Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -28,16 +27,14 @@ class Autoencoder(torch.nn.Module):
     and how they operate.
     """
 
-    def __init__(
-        self, skip_connection: Optional[SkipConnection] = None
-    ) -> None:
+    def __init__(self, skip_connection: SkipConnection | None = None) -> None:
         super().__init__()
         self.skip_connection = skip_connection
         self.blocks = torch.nn.ModuleList()
 
     def encode(
         self, *X: Tensor, return_states: bool = False
-    ) -> Union[Tensor, Tuple[Tensor, Sequence]]:
+    ) -> Tensor | tuple[Tensor, Sequence]:
         states = []
         for block in self.blocks:
             if isinstance(X, tuple):
@@ -53,7 +50,7 @@ class Autoencoder(torch.nn.Module):
             return X, states[:-1]
         return X
 
-    def decode(self, *X, states: Optional[Sequence[Tensor]] = None) -> Tensor:
+    def decode(self, *X, states: Sequence[Tensor] | None = None) -> Tensor:
         if self.skip_connection is not None and states is None:
             raise ValueError(
                 "Must pass intermediate states when autoencoder "
@@ -62,9 +59,8 @@ class Autoencoder(torch.nn.Module):
         elif states is not None:
             if len(states) != len(self.blocks) - 1:
                 raise ValueError(
-                    "Passed {} intermediate states, expected {}".format(
-                        len(states), len(self.blocks) - 1
-                    )
+                    f"Passed {len(states)} intermediate states, expected "
+                    f"{len(self.blocks) - 1}"
                 )
 
             # Don't skip connect the output layer

ml4gw/nn/autoencoder/convolutional.py

@@ -1,5 +1,4 @@
 from collections.abc import Callable, Sequence
-from typing import Optional
 
 import torch
 from torch import Tensor
@@ -21,9 +20,9 @@ class ConvBlock(Autoencoder):
         groups: int = 1,
         activation: torch.nn.Module = torch.nn.ReLU,
         norm: Module = torch.nn.BatchNorm1d,
-        decode_channels: Optional[int] = None,
-        output_activation: Optional[torch.nn.Module] = None,
-        skip_connection: Optional[SkipConnection] = None,
+        decode_channels: int | None = None,
+        output_activation: torch.nn.Module | None = None,
+        skip_connection: SkipConnection | None = None,
     ) -> None:
         super().__init__(skip_connection=None)
 
@@ -98,10 +97,10 @@ class ConvolutionalAutoencoder(Autoencoder):
         stride: int = 1,
         groups: int = 1,
         activation: torch.nn.Module = torch.nn.ReLU,
-        output_activation: Optional[torch.nn.Module] = None,
+        output_activation: torch.nn.Module | None = None,
         norm: Module = torch.nn.BatchNorm1d,
-        decode_channels: Optional[int] = None,
-        skip_connection: Optional[SkipConnection] = None,
+        decode_channels: int | None = None,
+        skip_connection: SkipConnection | None = None,
     ) -> None:
         # TODO: how to do this dynamically? Maybe the base
         # architecture looks for overlapping arguments between
@@ -145,9 +144,7 @@ class ConvolutionalAutoencoder(Autoencoder):
             self.blocks.append(block)
             in_channels = channels * groups
 
-    def decode(
-        self, *X, states=None, input_size: Optional[int] = None
-    ) -> Tensor:
+    def decode(self, *X, states=None, input_size: int | None = None) -> Tensor:
         X = super().decode(*X, states=states)
         if input_size is not None:
             return match_size(X, input_size)

ml4gw/nn/autoencoder/skip_connection.py

@@ -35,13 +35,11 @@ class ConcatSkipConnect(SkipConnection):
         rem = num_channels % self.groups
         if rem:
             raise ValueError(
-                "Number of channels in input tensor {} cannot "
-                "be divided evenly into {} groups".format(
-                    num_channels, self.groups
-                )
+                f"Number of channels in input tensor {num_channels} cannot "
+                f"be divided evenly into {self.groups} groups"
             )
 
         X = torch.split(X, self.groups, dim=1)
         state = torch.split(state, self.groups, dim=1)
-        frags = [i for j in zip(X, state) for i in j]
+        frags = [i for j in zip(X, state, strict=True) for i in j]
         return torch.cat(frags, dim=1)

ml4gw/nn/norm.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from collections.abc import Callable
 
 import torch
 from jaxtyping import Float
@@ -16,7 +16,7 @@ class GroupNorm1D(torch.nn.Module):
     def __init__(
         self,
         num_channels: int,
-        num_groups: Optional[int] = None,
+        num_groups: int | None = None,
         eps: float = 1e-5,
     ):
         super().__init__()
@@ -77,7 +77,7 @@ class GroupNorm1DGetter:
     for command-line parameterization with jsonargparse.
     """
 
-    def __init__(self, groups: Optional[int] = None) -> None:
+    def __init__(self, groups: int | None = None) -> None:
         self.groups = groups
 
     def __call__(self, num_channels: int) -> torch.nn.Module:
@@ -96,7 +96,7 @@ class GroupNorm2DGetter:
     for command-line parameterization with jsonargparse.
     """
 
-    def __init__(self, groups: Optional[int] = None) -> None:
+    def __init__(self, groups: int | None = None) -> None:
         self.groups = groups
 
     def __call__(self, num_channels: int) -> torch.nn.Module:

ml4gw/nn/resnet/resnet_1d.py

@@ -7,7 +7,8 @@ where training-time statistics are entirely arbitrary due to
 simulations.
 """
 
-from typing import Callable, List, Literal, Optional
+from collections.abc import Callable
+from typing import Literal
 
 import torch
 import torch.nn as nn
@@ -58,11 +59,11 @@ class BasicBlock(nn.Module):
         planes: int,
         kernel_size: int = 3,
         stride: int = 1,
-        downsample: Optional[nn.Module] = None,
+        downsample: nn.Module | None = None,
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Callable[..., nn.Module] | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -123,11 +124,11 @@ class Bottleneck(nn.Module):
         planes: int,
         kernel_size: int = 3,
         stride: int = 1,
-        downsample: Optional[nn.Module] = None,
+        downsample: nn.Module | None = None,
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: Optional[NormLayer] = None,
+        norm_layer: NormLayer | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -231,14 +232,14 @@ class ResNet1D(nn.Module):
     def __init__(
         self,
         in_channels: int,
-        layers: List[int],
+        layers: list[int],
         classes: int,
         kernel_size: int = 3,
         zero_init_residual: bool = False,
         groups: int = 1,
         width_per_group: int = 64,
-        stride_type: Optional[List[Literal["stride", "dilation"]]] = None,
-        norm_layer: Optional[NormLayer] = None,
+        stride_type: list[Literal["stride", "dilation"]] | None = None,
+        norm_layer: NormLayer | None = None,
     ) -> None:
         super().__init__()
 
@@ -257,10 +258,8 @@ class ResNet1D(nn.Module):
             stride_type = ["stride"] * (len(layers) - 1)
         if len(stride_type) != (len(layers) - 1):
             raise ValueError(
-                (
-                    "'stride_type' should be None or a {}-element "
-                    "tuple, got {}"
-                ).format(len(layers) - 1, stride_type)
+                f"'stride_type' should be None or a {len(layers) - 1}-element "
+                f"tuple, got {stride_type}"
             )
 
         self.groups = groups
@@ -289,7 +288,7 @@ class ResNet1D(nn.Module):
         # striding or dilating depending on the stride_type
         # argument)
         residual_layers = [self._make_layer(64, layers[0], kernel_size)]
-        it = zip(layers[1:], stride_type)
+        it = zip(layers[1:], stride_type, strict=True)
         for i, (num_blocks, stride) in enumerate(it):
             block_size = 64 * 2 ** (i + 1)
             layer = self._make_layer(

ml4gw/nn/resnet/resnet_2d.py

@@ -4,7 +4,8 @@ https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py
 but with arbitrary kernel sizes
 """
 
-from typing import Callable, List, Literal, Optional
+from collections.abc import Callable
+from typing import Literal
 
 import torch
 import torch.nn as nn
@@ -55,11 +56,11 @@ class BasicBlock(nn.Module):
         planes: int,
         kernel_size: int = 3,
         stride: int = 1,
-        downsample: Optional[nn.Module] = None,
+        downsample: nn.Module | None = None,
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Callable[..., nn.Module] | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -120,11 +121,11 @@ class Bottleneck(nn.Module):
         planes: int,
         kernel_size: int = 3,
         stride: int = 1,
-        downsample: Optional[nn.Module] = None,
+        downsample: nn.Module | None = None,
         groups: int = 1,
         base_width: int = 64,
         dilation: int = 1,
-        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Callable[..., nn.Module] | None = None,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -232,14 +233,14 @@ class ResNet2D(nn.Module):
     def __init__(
         self,
         in_channels: int,
-        layers: List[int],
+        layers: list[int],
         classes: int,
         kernel_size: int = 3,
         zero_init_residual: bool = False,
         groups: int = 1,
         width_per_group: int = 64,
-        stride_type: Optional[List[Literal["stride", "dilation"]]] = None,
-        norm_layer: Optional[NormLayer] = None,
+        stride_type: list[Literal["stride", "dilation"]] | None = None,
+        norm_layer: NormLayer | None = None,
     ) -> None:
         super().__init__()
         # default to using InstanceNorm if no
@@ -257,10 +258,8 @@ class ResNet2D(nn.Module):
             stride_type = ["stride"] * (len(layers) - 1)
         if len(stride_type) != (len(layers) - 1):
             raise ValueError(
-                (
-                    "'stride_type' should be None or a {}-element "
-                    "tuple, got {}"
-                ).format(len(layers) - 1, stride_type)
+                f"'stride_type' should be None or a {len(layers) - 1}-element "
+                f"tuple, got {stride_type}"
             )
 
         self.groups = groups
@@ -289,7 +288,7 @@ class ResNet2D(nn.Module):
         # striding or dilating depending on the stride_type
         # argument)
         residual_layers = [self._make_layer(64, layers[0], kernel_size)]
-        it = zip(layers[1:], stride_type)
+        it = zip(layers[1:], stride_type, strict=True)
         for i, (num_blocks, stride) in enumerate(it):
             block_size = 64 * 2 ** (i + 1)
             layer = self._make_layer(
@@ -316,7 +315,7 @@ class ResNet2D(nn.Module):
                 nn.init.kaiming_normal_(
                     m.weight, mode="fan_out", nonlinearity="relu"
                 )
-            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+            elif isinstance(m, nn.BatchNorm2d | nn.GroupNorm):
                 nn.init.constant_(m.weight, 1)
                 nn.init.constant_(m.bias, 0)
 

ml4gw/nn/streaming/online_average.py

@@ -1,5 +1,3 @@
-from typing import Optional, Tuple
-
 import torch
 from jaxtyping import Float
 from torch import Tensor
@@ -38,7 +36,7 @@ class OnlineAverager(torch.nn.Module):
         batch_size: int,
         num_updates: int,
         num_channels: int,
-        offset: Optional[int] = None,
+        offset: int | None = None,
     ) -> None:
         super().__init__()
         self.update_size = update_size
@@ -76,8 +74,8 @@ class OnlineAverager(torch.nn.Module):
     def forward(
         self,
         update: Float[Tensor, "batch channel time1"],
-        state: Optional[Float[Tensor, "channel time2"]] = None,
-    ) -> Tuple[Float[Tensor, "channel time3"], Float[Tensor, "channel time4"]]:
+        state: Float[Tensor, "channel time2"] | None = None,
+    ) -> tuple[Float[Tensor, "channel time3"], Float[Tensor, "channel time4"]]:
         if state is None:
             state = self.get_initial_state()
 

ml4gw/nn/streaming/snapshotter.py

@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, Tuple
+from collections.abc import Sequence
 
 import torch
 from jaxtyping import Float
@@ -58,15 +58,13 @@ class Snapshotter(torch.nn.Module):
         snapshot_size: int,
         stride_size: int,
         batch_size: int,
-        channels_per_snapshot: Optional[Sequence[int]] = None,
+        channels_per_snapshot: Sequence[int] | None = None,
     ) -> None:
         super().__init__()
         if stride_size >= snapshot_size:
             raise ValueError(
-                "Snapshotter can't accommodate stride {} "
-                "which is greater than snapshot size {}".format(
-                    stride_size, snapshot_size
-                )
+                f"Snapshotter can't accommodate stride {stride_size} "
+                f"which is greater than snapshot size {snapshot_size}"
             )
 
         self.snapshot_size = snapshot_size
@@ -77,9 +75,8 @@ class Snapshotter(torch.nn.Module):
         if channels_per_snapshot is not None:
             if sum(channels_per_snapshot) != num_channels:
                 raise ValueError(
-                    "Can't break {} channels into {}".format(
-                        num_channels, channels_per_snapshot
-                    )
+                    f"Can't break {num_channels} channels into "
+                    f"{channels_per_snapshot}"
                 )
         self.channels_per_snapshot = channels_per_snapshot
         self.num_channels = num_channels
@@ -90,8 +87,8 @@ class Snapshotter(torch.nn.Module):
     def forward(
         self,
         update: Float[Tensor, "channel time1"],
-        snapshot: Optional[Float[Tensor, "channel time2"]] = None,
-    ) -> Tuple[Tensor, ...]:
+        snapshot: Float[Tensor, "channel time2"] | None = None,
+    ) -> tuple[Tensor, ...]:
         if snapshot is None:
             snapshot = self.get_initial_state()
 
@@ -108,9 +105,8 @@ class Snapshotter(torch.nn.Module):
         if self.channels_per_snapshot is not None:
             if snapshots.size(1) != self.num_channels:
                 raise ValueError(
-                    "Expected {} channels, found {}".format(
-                        self.num_channels, snapshots.size(1)
-                    )
+                    f"Expected {self.num_channels} channels, found "
+                    f"{snapshots.size(1)}"
                 )
             snapshots = torch.split(
                 snapshots, self.channels_per_snapshot, dim=1