pytme 0.2.9__cp311-cp311-macosx_15_0_arm64.whl → 0.3.0__cp311-cp311-macosx_15_0_arm64.whl

tme/analyzer/proxy.py

@@ -0,0 +1,123 @@
+"""
+Implements SharedAnalyzerProxy to managed shared memory of Analyzer instances
+across different tasks.
+
+This is primarily useful for CPU template matching, where parallelization can
+be performed over rotations, rather than subsections of a large input volume.
+
+Copyright (c) 2025 European Molecular Biology Laboratory
+
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple
+from multiprocessing import Manager
+from multiprocessing.shared_memory import SharedMemory
+
+from ..backends import backend as be
+
+__all__ = ["StatelessSharedAnalyzerProxy", "SharedAnalyzerProxy"]
+
+
+class StatelessSharedAnalyzerProxy:
+    """
+    Proxy that wraps functional analyzers for concurrent access via shared memory.
+
+    Enables multiple processes/threads to safely update the same analyzer
+    while preserving the functional interface of the underlying analyzer.
+    """
+
+    def __init__(self, analyzer_class: type, analyzer_params: dict):
+        self._shared = False
+        self._process = self._direct_call
+
+        self._analyzer = analyzer_class(**analyzer_params)
+
+    def __call__(self, state, *args, **kwargs):
+        return self._process(state, *args, **kwargs)
+
+    def init_state(self, shm_handler=None, *args, **kwargs) -> Tuple:
+        state = self._analyzer.init_state()
+        if shm_handler is not None:
+            self._shared = True
+            state = self._to_shared(state, shm_handler)
+
+            self._lock = Manager().Lock()
+            self._process = self._thread_safe_call
+        return state
+
+    def _to_shared(self, state: Tuple, shm_handler):
+        backend_arr = type(be.zeros((1), dtype=be._float_dtype))
+
+        ret = []
+        for v in state:
+            if isinstance(v, backend_arr):
+                v = be.to_sharedarr(v, shm_handler)
+            elif isinstance(v, dict):
+                v = Manager().dict(**v)
+            ret.append(v)
+        return tuple(ret)
+
+    def _shared_to_object(self, shared: type):
+        if not self._shared:
+            return shared
+        if isinstance(shared, tuple) and len(shared):
+            if isinstance(shared[0], SharedMemory):
+                return be.from_sharedarr(shared)
+        return shared
+
+    def _thread_safe_call(self, state, *args, **kwargs):
+        """Thread-safe call to analyzer"""
+        with self._lock:
+            state = tuple(self._shared_to_object(x) for x in state)
+            return self._direct_call(state, *args, **kwargs)
+
+    def _direct_call(self, state, *args, **kwargs):
+        """Direct call to analyzer without locking"""
+        return self._analyzer(state, *args, **kwargs)
+
+    def result(self, state, **kwargs):
+        """Extract final result"""
+        final_state = state
+        if self._shared:
+            # Convert shared arrays back to regular arrays and copy to
+            # avoid array invalidation by shared memory handler
+            final_state = tuple(self._shared_to_object(x) for x in final_state)
+        return self._analyzer.result(final_state, **kwargs)
+
+    def merge(self, *args, **kwargs):
+        return self._analyzer.merge(*args, **kwargs)
+
+
+class SharedAnalyzerProxy(StatelessSharedAnalyzerProxy):
+    """
+    Child of :py:class:`StatelessSharedAnalyzerProxy` that is aware
+    of the current analyzer state to emulate the previous analyzer interface.
+    """
+
+    def __init__(
+        self,
+        analyzer_class: type,
+        analyzer_params: dict,
+        shm_handler: type = None,
+        **kwargs,
+    ):
+        super().__init__(
+            analyzer_class=analyzer_class,
+            analyzer_params=analyzer_params,
+        )
+        if not self._analyzer.shareable:
+            shm_handler = None
+        self.init_state(shm_handler)
+
+    def init_state(self, shm_handler=None, *args, **kwargs) -> Tuple:
+        self._state = super().init_state(shm_handler, *args, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        state = super().__call__(self._state, *args, **kwargs)
+        if not self._shared:
+            self._state = state
+
+    def result(self, **kwargs):
+        """Extract final result"""
+        return super().result(self._state, **kwargs)

tme/backends/__init__.py

@@ -1,8 +1,9 @@
-""" pyTME backend manager.
+"""
+pyTME backend manager.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Dict, List

tme/backends/_cupy_utils.py

@@ -1,34 +1,35 @@
-""" Utility functions for cupy backend.
+"""
+Utility functions for cupy backend.
 
-    The functions spline_filter, _prepad_for_spline_filter, _filter_input,
-    _get_coord_affine_batched and affine_transform are largely copied from
-    cupyx.scipy.ndimage which operates under the following license
+The functions spline_filter, _prepad_for_spline_filter, _filter_input,
+_get_coord_affine_batched and affine_transform are largely copied from
+cupyx.scipy.ndimage which operates under the following license
 
-    Copyright (c) 2015 Preferred Infrastructure, Inc.
-    Copyright (c) 2015 Preferred Networks, Inc.
+Copyright (c) 2015 Preferred Infrastructure, Inc.
+Copyright (c) 2015 Preferred Networks, Inc.
 
-    Permission is hereby granted, free of charge, to any person obtaining a copy
-    of this software and associated documentation files (the "Software"), to deal
-    in the Software without restriction, including without limitation the rights
-    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-    copies of the Software, and to permit persons to whom the Software is
-    furnished to do so, subject to the following conditions:
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
 
-    The above copyright notice and this permission notice shall be included in
-    all copies or substantial portions of the Software.
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
 
-    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-    THE SOFTWARE.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
 
-    I have since extended the functionality of the cupyx.scipy.ndimage functions
-    in question to support batched inputs.
+I have since extended the functionality of the cupyx.scipy.ndimage functions
+in question to support batched inputs.
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import numpy

tme/backends/_jax_utils.py

@@ -1,24 +1,28 @@
-""" Utility functions for jax backend.
+"""
+Utility functions for jax backend.
 
-    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+Copyright (c) 2023-2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Tuple
 from functools import partial
 
 import jax.numpy as jnp
-from jax import pmap, lax
+from jax import pmap, lax, vmap
 
 from ..types import BackendArray
 from ..backends import backend as be
 from ..matching_utils import normalize_template as _normalize_template
 
 
+__all__ = ["scan"]
+
+
 def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
     """
-    Computes :py:meth:`tme.matching_exhaustive.cc_setup`.
+    Computes :py:meth:`tme.matching_scores.cc_setup`.
     """
     template_ft = jnp.fft.rfftn(template, s=template.shape)
     template_ft = template_ft.at[:].multiply(ft_target)
@@ -27,18 +31,17 @@ def _correlate(template: BackendArray, ft_target: BackendArray) -> BackendArray:
 
 
 def _flc_scoring(
-    template: BackendArray,
-    template_mask: BackendArray,
     ft_target: BackendArray,
     ft_target2: BackendArray,
+    template: BackendArray,
+    template_mask: BackendArray,
     n_observations: BackendArray,
     eps: float,
     **kwargs,
 ) -> BackendArray:
     """
-    Computes :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    Computes :py:meth:`tme.matching_scores.flc_scoring`.
     """
-    correlation = _correlate(template=template, ft_target=ft_target)
     inv_denominator = _reciprocal_target_std(
         ft_target=ft_target,
         ft_target2=ft_target2,
@@ -46,18 +49,17 @@ def _flc_scoring(
         eps=eps,
         n_observations=n_observations,
     )
-    correlation = correlation.at[:].multiply(inv_denominator)
-    return correlation
+    return _flcSphere_scoring(ft_target, template, inv_denominator)
 
 
 def _flcSphere_scoring(
-    template: BackendArray,
     ft_target: BackendArray,
+    template: BackendArray,
     inv_denominator: BackendArray,
     **kwargs,
 ) -> BackendArray:
     """
-    Computes :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    Computes :py:meth:`tme.matching_scores.corr_scoring`.
     """
     correlation = _correlate(template=template, ft_target=ft_target)
     correlation = correlation.at[:].multiply(inv_denominator)
@@ -76,7 +78,7 @@ def _reciprocal_target_std(
 
     See Also
     --------
-    :py:meth:`tme.matching_exhaustive.flc_scoring`.
+    :py:meth:`tme.matching_scores.flc_scoring`.
     """
     ft_shape = template_mask.shape
     ft_template_mask = jnp.fft.rfftn(template_mask, s=ft_shape)
@@ -162,12 +164,12 @@ def scan(
         template_rot = _normalize_template(
             template_rot, template_mask_rot, n_observations
         )
-        template_rot = be.topleft_pad(template_rot, fast_shape)
-        template_mask_rot = be.topleft_pad(template_mask_rot, fast_shape)
+        rot_pad = be.topleft_pad(template_rot, fast_shape)
+        mask_rot_pad = be.topleft_pad(template_mask_rot, fast_shape)
 
         scores = scoring_func(
-            template=template_rot,
-            template_mask=template_mask_rot,
+            template=rot_pad,
+            template_mask=mask_rot_pad,
             ft_target=ft_target,
             ft_target2=ft_target2,
             inv_denominator=inv_denominator,

tme/backends/cupy_backend.py

@@ -1,17 +1,15 @@
-""" Backend using cupy for template matching.
+"""
+Backend using cupy for template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-import warnings
 from importlib.util import find_spec
 from contextlib import contextmanager
 from typing import Tuple, Callable, List
 
-import numpy as np
-
 from .npfftw_backend import NumpyFFTWBackend
 from ..types import CupyArray, NDArray, shm_type
 
@@ -113,16 +111,6 @@ class CupyBackend(NumpyFFTWBackend):
     def unravel_index(self, indices, shape):
         return self._array_backend.unravel_index(indices=indices, dims=shape)
 
-    def unique(self, ar, axis=None, *args, **kwargs):
-        if axis is None:
-            return self._array_backend.unique(ar=ar, axis=axis, *args, **kwargs)
-
-        warnings.warn("Axis argument not yet supported in CupY, falling back to NumPy.")
-        ret = np.unique(ar=self.to_numpy_array(ar), axis=axis, *args, **kwargs)
-        if not isinstance(ret, tuple):
-            return self.to_backend_array(ret)
-        return tuple(self.to_backend_array(k) for k in ret)
-
     def build_fft(
         self,
         fwd_shape: Tuple[int],
@@ -155,15 +143,14 @@ class CupyBackend(NumpyFFTWBackend):
         def rfftn(
             arr: CupyArray, out: CupyArray = None, s=rfft_shape, axes=fwd_axes
         ) -> CupyArray:
-            return self.rfftn(arr, s=s, axes=fwd_axes)
+            return self.rfftn(arr, s=s, axes=fwd_axes, overwrite_x=True)
 
         def irfftn(
             arr: CupyArray, out: CupyArray = None, s=irfft_shape, axes=inv_axes
         ) -> CupyArray:
-            return self.irfftn(arr, s=s, axes=inv_axes)
+            return self.irfftn(arr, s=s, axes=inv_axes, overwrite_x=True)
 
         PLAN_CACHE[current_device] = [fwd_shape, inv_shape]
-
         return rfftn, irfftn
 
     def rfftn(self, arr: CupyArray, out: CupyArray = None, **kwargs) -> CupyArray:
@@ -248,13 +235,13 @@ class CupyBackend(NumpyFFTWBackend):
             )
             return None
 
-        # if data.ndim == 3 and cache and self.texture_available:
-        #     # Device memory pool (should) come to rescue performance
-        #     temp = self.zeros(data.shape, data.dtype)
-        #     texture = self._get_texture(data, order=order, prefilter=prefilter)
-        #     texture.affine(transform_m=matrix, profile=False, output=temp)
-        #     output[out_slice] = temp
-        #     return None
+        if data.ndim == 3 and cache and self.texture_available:
+            # Device memory pool (should) come to rescue performance
+            temp = self.zeros(data.shape, data.dtype)
+            texture = self._get_texture(data, order=order, prefilter=prefilter)
+            texture.affine(transform_m=matrix, profile=False, output=temp)
+            output[out_slice] = temp
+            return None
 
         self.affine_transform(
             input=data,

tme/backends/jax_backend.py

@@ -1,12 +1,13 @@
-""" Backend using jax for template matching.
+"""
+Backend using jax for template matching.
 
-    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+Copyright (c) 2023-2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from functools import wraps
-from typing import Tuple, List, Callable
+from typing import Tuple, List, Callable, Dict
 
 from ..types import BackendArray
 from .npfftw_backend import NumpyFFTWBackend, shm_type
@@ -50,12 +51,6 @@ class JaxBackend(NumpyFFTWBackend):
         )
         self.scipy = jsp
         self._create_ufuncs()
-        try:
-            from ._jax_utils import scan as _
-
-            self.scan = self._scan
-        except Exception:
-            pass
 
     def from_sharedarr(self, arr: BackendArray) -> BackendArray:
         return arr
@@ -188,7 +183,18 @@ class JaxBackend(NumpyFFTWBackend):
         rotations = rotations.at[:].set(self.where(update, rotations, rotation_index))
         return max_scores, rotations
 
-    def _scan(
+    def compute_convolution_shapes(
+        self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
+    ) -> Tuple[List[int], List[int], List[int]]:
+        from scipy.fft import next_fast_len
+
+        convolution_shape = [int(x + y - 1) for x, y in zip(arr1_shape, arr2_shape)]
+        fast_shape = [next_fast_len(x, real=True) for x in convolution_shape]
+        fast_ft_shape = list(fast_shape[:-1]) + [fast_shape[-1] // 2 + 1]
+
+        return convolution_shape, fast_shape, fast_ft_shape
+
+    def scan(
         self,
         matching_data: type,
         splits: Tuple[Tuple[slice, slice]],
@@ -213,9 +219,9 @@ class JaxBackend(NumpyFFTWBackend):
         conv_shape, fast_shape, fast_ft_shape, shift = matching_data._fourier_padding(
             target_shape=self.to_numpy_array(target_shape),
             template_shape=self.to_numpy_array(matching_data._template.shape),
-            pad_fourier=False,
+            batch_mask=self.to_numpy_array(matching_data._batch_mask),
+            pad_target=pad_target,
         )
-
         analyzer_args = {
             "convolution_mode": convolution_mode,
             "fourier_shift": shift,
@@ -245,19 +251,18 @@ class JaxBackend(NumpyFFTWBackend):
 
             targets, translation_offsets = [], []
             for target_split, template_split in split_subset:
-                base = matching_data.subset_by_slice(
+                base, translation_offset = matching_data.subset_by_slice(
                     target_slice=target_split,
                     target_pad=target_pad,
                     template_slice=template_split,
                 )
-                translation_offsets.append(base._translation_offset)
+                translation_offsets.append(translation_offset)
                 targets.append(self.topleft_pad(base._target, fast_shape))
 
             if create_filter:
                 filter_args = {
                     "data_rfft": self.fft.rfftn(targets[0]),
                     "return_real_fourier": True,
-                    "shape_is_real_fourier": False,
                 }
 
             if create_template_filter:
@@ -287,15 +292,11 @@ class JaxBackend(NumpyFFTWBackend):
 
             for index in range(scores.shape[0]):
                 temp = callback_class(
-                    shape=scores.shape,
-                    scores=scores[index],
-                    rotations=rotations[index],
-                    thread_safe=False,
+                    shape=scores[index].shape,
                     offset=translation_offsets[index],
                 )
-                temp.rotation_mapping = rotation_mapping
-                ret.append(tuple(temp._postprocess(**analyzer_args)))
-
+                state = (scores[index], rotations[index], rotation_mapping)
+                ret.append(temp.result(state, **analyzer_args))
         return ret
 
     def get_available_memory(self) -> int:

tme/backends/matching_backend.py

@@ -1,8 +1,9 @@
-""" Strategy pattern to allow for flexible array / FFT backends.
+"""
+Strategy pattern to allow for flexible array / FFT backends.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from abc import ABC, abstractmethod

tme/backends/mlx_backend.py

@@ -1,8 +1,9 @@
-""" Backend using Apple's MLX library for template matching.
+"""
+Backend using Apple's MLX library for template matching.
 
-    Copyright (c) 2024 European Molecular Biology Laboratory
+Copyright (c) 2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Tuple, List, Callable

tme/backends/npfftw_backend.py

@@ -1,8 +1,9 @@
-""" Backend using numpy and pyFFTW for template matching.
+"""
+Backend using numpy and pyFFTW for template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import os
@@ -159,8 +160,9 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         temp = self._array_backend.zeros(1, dtype=dtype)
         return temp.nbytes
 
-    @staticmethod
-    def astype(arr, dtype: Type) -> NDArray:
+    def astype(self, arr, dtype: Type) -> NDArray:
+        if self._array_backend.iscomplexobj(arr):
+            arr = arr.real
         return arr.astype(dtype)
 
     @staticmethod
@@ -396,33 +398,33 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         out_mask: NDArray = None,
         order: int = 3,
         cache: bool = False,
+        batched: bool = False,
     ) -> Tuple[NDArray, NDArray]:
-        out = self.zeros_like(arr) if out is None else out
-        batched = arr.ndim != rotation_matrix.shape[0]
-
-        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
-        if not use_geometric_center:
-            center = self.center_of_mass(arr, cutoff=0)
-
-        offset = int(arr.ndim - rotation_matrix.shape[0])
-        center = center[offset:]
-        translation = self.zeros(center.size) if translation is None else translation
-        matrix = self._rigid_transform_matrix(
-            rotation_matrix=rotation_matrix,
-            translation=translation,
-            center=center,
-        )
-
-        subset = tuple(slice(None) for _ in range(arr.ndim))
-        if offset > 1:
-            subset = tuple(
-                0 if i < (offset - 1) else slice(None) for i in range(arr.ndim)
+        if out is None:
+            out = self.zeros_like(arr)
+
+        # Check whether rotation_matrix is already a rigid transform matrix
+        matrix = rotation_matrix
+        if matrix.shape[-1] == (arr.ndim - int(batched)):
+            center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
+            if not use_geometric_center:
+                center = self.center_of_mass(arr, cutoff=0)
+
+            offset = int(arr.ndim - rotation_matrix.shape[0])
+            center = center[offset:]
+            translation = (
+                self.zeros(center.size) if translation is None else translation
+            )
+            matrix = self._rigid_transform_matrix(
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                center=center,
             )
 
         self._rigid_transform(
-            data=arr[subset],
+            data=arr,
             matrix=matrix,
-            output=out[subset],
+            output=out,
             order=order,
             prefilter=True,
             cache=cache,
@@ -431,11 +433,13 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
 
         # Applying the prefilter leads to artifacts in the mask.
         if arr_mask is not None:
-            out_mask = self.zeros_like(arr_mask) if out_mask is None else out_mask
+            if out_mask is None:
+                out_mask = self.zeros_like(arr_mask)
+
             self._rigid_transform(
-                data=arr_mask[subset],
+                data=arr_mask,
                 matrix=matrix,
-                output=out_mask[subset],
+                output=out_mask,
                 order=order,
                 prefilter=False,
                 cache=cache,

tme/backends/pytorch_backend.py

@@ -1,9 +1,10 @@
-""" Backend using pytorch and optionally GPU acceleration for
-    template matching.
+"""
+Backend using pytorch and optionally GPU acceleration for
+template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Tuple, Callable
@@ -134,6 +135,15 @@ class PytorchBackend(NumpyFFTWBackend):
     def astype(self, arr: TorchTensor, dtype: type) -> TorchTensor:
         return arr.to(dtype)
 
+    @staticmethod
+    def at(arr, idx, value) -> NDArray:
+        arr[idx] = value
+        return arr
+
+    @staticmethod
+    def addat(arr, indices, *args, **kwargs) -> NDArray:
+        return arr.index_put_(indices, *args, accumulate=True, **kwargs)
+
     def flip(self, a, axis, **kwargs):
         return self._array_backend.flip(input=a, dims=axis, **kwargs)
 
@@ -296,6 +306,9 @@ class PytorchBackend(NumpyFFTWBackend):
         kwargs["dim"] = kwargs.pop("axes", None)
         return self._array_backend.fft.irfftn(arr, **kwargs)
 
+    def _rigid_transform_matrix(self, rotation_matrix, *args, **kwargs):
+        return rotation_matrix
+
     def rigid_transform(
         self,
         arr: TorchTensor,
@@ -307,6 +320,7 @@ class PytorchBackend(NumpyFFTWBackend):
         out_mask: TorchTensor = None,
         order: int = 1,
         cache: bool = False,
+        **kwargs,
     ):
         _mode_mapping = {0: "nearest", 1: "bilinear", 3: "bicubic"}
         mode = _mode_mapping.get(order, None)