From 63b69429ce547ed1b2fe24ea0945a366258f96f7 Mon Sep 17 00:00:00 2001 From: Olivier Verdier Date: Sun, 29 Sep 2024 15:56:10 +0200 Subject: [PATCH] MNT: compat Python 3.9 --- diffeopt/group/ddmatch/action/density.py | 6 +++--- diffeopt/group/ddmatch/action/function.py | 8 ++++---- diffeopt/group/ddmatch/representation.py | 6 +++--- diffeopt/group/representation.py | 4 ++-- diffeopt/optim.py | 6 +++--- 5 files changed, 15 insertions(+), 15 deletions(-) diff --git a/diffeopt/group/ddmatch/action/density.py b/diffeopt/group/ddmatch/action/density.py index d170ac6..ed46b29 100644 --- a/diffeopt/group/ddmatch/action/density.py +++ b/diffeopt/group/ddmatch/action/density.py @@ -1,4 +1,4 @@ -from typing import Callable, Any +from typing import Callable, Any, Union import torch import numpy as np from ddmatch.core import ( # type: ignore @@ -10,7 +10,7 @@ # TODO: implement forward or backward in numba?? -def get_density_action(shape: torch.Size, compute_id: bool=False) -> Callable[[torch.Tensor, torch.Tensor | ForwardDiffeo], torch.Tensor]: +def get_density_action(shape: torch.Size, compute_id: bool=False) -> Callable[[torch.Tensor, Union[torch.Tensor, ForwardDiffeo]], torch.Tensor]: image = np.zeros(shape) compute_grad = generate_optimized_image_gradient(image) compute_pullback = generate_optimized_image_composition(image) @@ -20,7 +20,7 @@ class DensityAction(torch.autograd.Function): """ @staticmethod - def forward(ctx: Any, x: torch.Tensor, g: torch.Tensor | ForwardDiffeo) -> torch.Tensor: + def forward(ctx: Any, x: torch.Tensor, g: Union[torch.Tensor, ForwardDiffeo]) -> torch.Tensor: ctx.save_for_backward(x) if isinstance(g, ForwardDiffeo): torch_data = g.forward diff --git a/diffeopt/group/ddmatch/action/function.py b/diffeopt/group/ddmatch/action/function.py index a891456..0504496 100644 --- a/diffeopt/group/ddmatch/action/function.py +++ b/diffeopt/group/ddmatch/action/function.py @@ -1,4 +1,4 @@ -from typing import Any, Callable +from typing import Any, Callable, Union, Optional import torch import numpy as np @@ -8,7 +8,7 @@ from ..group import DiffeoGroup -def get_composition_action(shape: torch.Size, compute_id:bool=False) -> Callable[[torch.Tensor, torch.Tensor | Diffeo], torch.Tensor]: +def get_composition_action(shape: torch.Size, compute_id:bool=False) -> Callable[[torch.Tensor, Union[torch.Tensor, Diffeo]], torch.Tensor]: """ compute_id: right composition with the identity is the identity map, but it takes time to compute. @@ -32,7 +32,7 @@ class CompositionAction(torch.autograd.Function): and its derivative wrt g. """ @staticmethod - def forward(ctx: Any, q: torch.Tensor, g: torch.Tensor | Diffeo) -> torch.Tensor: + def forward(ctx: Any, q: torch.Tensor, g: Union[torch.Tensor, Diffeo]) -> torch.Tensor: if isinstance(g, Diffeo): torch_data = g.forward to_save = g.backward @@ -52,7 +52,7 @@ def forward(ctx: Any, q: torch.Tensor, g: torch.Tensor | Diffeo) -> torch.Tensor return res @staticmethod - def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: # type: ignore[override] + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]]: # type: ignore[override] """ This is the adjoint of the derivative only if g was the identity. diff --git a/diffeopt/group/ddmatch/representation.py b/diffeopt/group/ddmatch/representation.py index 3316cf3..e490c81 100644 --- a/diffeopt/group/ddmatch/representation.py +++ b/diffeopt/group/ddmatch/representation.py @@ -1,4 +1,4 @@ -from typing import Callable +from typing import Callable, Union import torch from ..representation import Representation @@ -7,11 +7,11 @@ from .action.function import get_composition_action class FunctionRepresentation(Representation): - def get_representation(self, group: BaseDiffeoGroup) -> Callable[[torch.Tensor, torch.Tensor | Diffeo], torch.Tensor]: + def get_representation(self, group: BaseDiffeoGroup) -> Callable[[torch.Tensor, Union[torch.Tensor, Diffeo]], torch.Tensor]: return get_composition_action(group.shape) from .action.density import get_density_action class DensityRepresentation(Representation): - def get_representation(self, group: BaseDiffeoGroup) -> Callable[[torch.Tensor, torch.Tensor | Diffeo], torch.Tensor]: + def get_representation(self, group: BaseDiffeoGroup) -> Callable[[torch.Tensor, Union[torch.Tensor, Diffeo]], torch.Tensor]: return get_density_action(group.shape) diff --git a/diffeopt/group/representation.py b/diffeopt/group/representation.py index f2917a5..c145434 100644 --- a/diffeopt/group/representation.py +++ b/diffeopt/group/representation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Callable +from typing import Callable, Union from .deformation import Deformation import torch from torch.nn.parameter import Parameter @@ -39,5 +39,5 @@ def forward(self, I: torch.Tensor) -> torch.Tensor: return self.representation(self.representation(I, self.perturbation.base.deformation), self.perturbation) @abstractmethod - def get_representation(self, group: BaseDiffeoGroup) -> Callable[[torch.Tensor, torch.Tensor | Diffeo], torch.Tensor]: + def get_representation(self, group: BaseDiffeoGroup) -> Callable[[torch.Tensor, Union[torch.Tensor, Diffeo]], torch.Tensor]: pass diff --git a/diffeopt/optim.py b/diffeopt/optim.py index fed853d..ac56f7f 100644 --- a/diffeopt/optim.py +++ b/diffeopt/optim.py @@ -1,4 +1,4 @@ -from typing import Callable, Any +from typing import Callable, Any, Optional, Union from abc import ABC, abstractmethod import torch from .group.representation import Perturbation @@ -8,8 +8,8 @@ class DiffeoOptimizer(Optimizer, ABC): @torch.no_grad() - def step(self, closure:Callable[[], float] | None = None) -> float | None: # type: ignore[override] - loss: float | None = None + def step(self, closure:Optional[Callable[[], float]] = None) -> Optional[float]: # type: ignore[override] + loss: Optional[float] = None if closure is not None: with torch.enable_grad(): loss = closure()