<gh_stars>1-10
# -*- coding: utf-8 -*-

from typing import Tuple

import torch
from torch import Tensor


class Conv2dNoRho(torch.nn.Conv2d):
    def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, padding=0, dilation=1, groups=1,
                 bias: bool = True, padding_mode: str = 'zeros'):
        super(Conv2dNoRho, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias,
                                          padding_mode)

    def forward(self, u: Tensor, s: Tensor) -> Tuple[Tensor, Tensor]:
        u = self._conv_forward(u, self.weight, self.bias)
        if self.bias is not None:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), torch.pow(self.bias, 2)))
        else:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), self.bias))
        return u, s


class StdConv2d(torch.nn.Conv2d):
    def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, padding=0, dilation=1, groups=1,
                 bias: bool = True, padding_mode: str = 'zeros'):
        super(StdConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias,
                                        padding_mode)

    def forward(self, s: Tensor) -> Tensor:
        if self.bias is not None:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), torch.pow(self.bias, 2)))
        else:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), self.bias))
        return s

# -*- coding: utf-8 -*-

from typing import Tuple

import torch
from torch import Tensor


class Conv2dNoRho(torch.nn.Conv2d):
    def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, padding=0, dilation=1, groups=1,
                 bias: bool = True, padding_mode: str = 'zeros'):
        super(Conv2dNoRho, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias,
                                          padding_mode)

    def forward(self, u: Tensor, s: Tensor) -> Tuple[Tensor, Tensor]:
        u = self._conv_forward(u, self.weight, self.bias)
        if self.bias is not None:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), torch.pow(self.bias, 2)))
        else:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), self.bias))
        return u, s


class StdConv2d(torch.nn.Conv2d):
    def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, padding=0, dilation=1, groups=1,
                 bias: bool = True, padding_mode: str = 'zeros'):
        super(StdConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias,
                                        padding_mode)

    def forward(self, s: Tensor) -> Tensor:
        if self.bias is not None:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), torch.pow(self.bias, 2)))
        else:
            s = torch.sqrt(self._conv_forward(torch.pow(s, 2), torch.pow(self.weight, 2), self.bias))
        return s
# -*- coding: utf-8 -*-