'''ResNet using PSG in PyTorch.

For Pre-activation ResNet, see 'preact_resnet.py'.

Reference:
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
    Deep Residual Learning for Image Recognition. arXiv:1512.03385
'''
from numpy.lib.arraysetops import isin
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np

from models.masked_psg_seed_conv import PredictiveSeedConv2d
from masked_layers import layers


# Fixed
NUM_BITS = 32
NUM_BITS_WEIGHT = 32
NUM_BITS_GRAD = None

BIPRECISION = False
PREDICTIVE_FORWARD = False

WRITER = None
WRITER_PREFIX_COUNTER = 0

# Tunable
PREDICTIVE_BACKWARD = True

MSB_BITS = 4
MSB_BITS_WEIGHT = 4
MSB_BITS_GRAD = 8

THRESHOLD = 0.0
SPARSIFY = False
SIGN = True



def conv1x1(in_planes, out_planes, stride=1, input_signed=True, predictive_forward=True, writer_prefix=""):
    "1x1 convolution with no padding"
    predictive_forward = PREDICTIVE_FORWARD and predictive_forward
    return PredictiveSeedConv2d(
        in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False,
        num_bits=NUM_BITS, num_bits_weight=NUM_BITS_WEIGHT, num_bits_grad=NUM_BITS_GRAD,
        biprecision=BIPRECISION, input_signed=input_signed,
        predictive_forward=predictive_forward, predictive_backward=PREDICTIVE_BACKWARD,
        msb_bits=MSB_BITS, msb_bits_weight=MSB_BITS_WEIGHT, msb_bits_grad=MSB_BITS_GRAD,
        threshold=THRESHOLD, sparsify=SPARSIFY, sign=SIGN,
        writer=WRITER, writer_prefix=writer_prefix)


def conv3x3(in_planes, out_planes, stride=1, input_signed=False, predictive_forward=True, writer_prefix=""):
    "3x3 convolution with padding"
    predictive_forward = PREDICTIVE_FORWARD and predictive_forward
    return PredictiveSeedConv2d(
        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False,
        num_bits=NUM_BITS, num_bits_weight=NUM_BITS_WEIGHT, num_bits_grad=NUM_BITS_GRAD,
        biprecision=BIPRECISION, input_signed=input_signed,
        predictive_forward=predictive_forward, predictive_backward=PREDICTIVE_BACKWARD,
        msb_bits=MSB_BITS, msb_bits_weight=MSB_BITS_WEIGHT, msb_bits_grad=MSB_BITS_GRAD,
        threshold=THRESHOLD, sparsify=SPARSIFY, sign=SIGN,
        writer=WRITER, writer_prefix=writer_prefix)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(in_planes, planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes, stride=1, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                # nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
                conv1x1(in_planes, self.expansion*planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None),
                nn.BatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_planes, planes, stride=1):
        super(Bottleneck, self).__init__()
        # self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.conv1 = conv1x1(in_planes, planes, stride=1, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn1 = nn.BatchNorm2d(planes)
        # self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.conv2 = conv3x3(planes, planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn2 = nn.BatchNorm2d(planes)
        # self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)
        self.conv3 = conv1x1(planes, self.expansion*planes, stride=1, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn3 = nn.BatchNorm2d(self.expansion*planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                # nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
                conv1x1(in_planes, self.expansion*planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None),
                nn.BatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, block, num_blocks, in_planes=64, num_classes=10, init_method='standard'):
        super(ResNet, self).__init__()
        self.in_planes = in_planes

        self.conv1 = conv3x3(3, self.in_planes, stride=1, input_signed=True, predictive_forward=False, writer_prefix=None)
        self.bn1 = nn.BatchNorm2d(self.in_planes)
        if self.in_planes ==  64:
            # self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
            self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
            self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
            self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
            self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
            self.linear = nn.Linear(512*block.expansion, num_classes)
            #self.linear = layers.Linear(512*block.expansion, num_classes)
        elif self.in_planes == 16:
            self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
            self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
            self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)
            self.layer4 = None
            self.linear = nn.Linear(64, num_classes)

        self.reset_conv_parameters(init_method)
        print('conv weights reset to {}'.format(init_method))

    def reset_parameters(self, module, init_method="kaiming_uniform") -> None:
        if init_method == "kaiming_constant_signed":
            fan = nn.init._calculate_correct_fan(module.weight, "fan_in")
            gain = nn.init.calculate_gain("relu")
            std = gain / math.sqrt(fan)
            with torch.no_grad():
                module.weight.data = module.weight.data.sign() * std
        elif init_method == "kaiming_constant_unsigned":
            fan = nn.init._calculate_correct_fan(module.weight, "fan_in")
            gain = nn.init.calculate_gain("relu")
            std = gain / math.sqrt(fan)
            with torch.no_grad():
                module.weight.data = torch.ones_like(module.weight.data) * std
        elif init_method == "kaiming_normal":
            nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu")
        elif init_method == "kaiming_uniform":
            nn.init.kaiming_uniform_(module.weight, mode="fan_in", nonlinearity="relu")
        elif init_method == "kaiming_laplace":
            fan = nn.init._calculate_correct_fan(module.weight, "fan_in")
            gain = nn.init.calculate_gain("relu")
            scale = gain / math.sqrt(2.0 * fan)
            with torch.no_grad():
                new_weight = np.random.laplace(loc=0.0, scale=scale, size=module.weight.shape)
                module.weight.data = module.weight.data.new_tensor(torch.from_numpy(new_weight).clone().detach())
        elif init_method == "xavier_normal":
            nn.init.xavier_normal_(module.weight)
        elif init_method == "xavier_constant":
            fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(module.weight)
            std = math.sqrt(2.0 / float(fan_in + fan_out))
            with torch.no_grad():
                module.weight.data = module.weight.data.sign() * std
        elif init_method == "standard":
            nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
        else:
            raise ValueError(f"{init_method} is not an initialization option!")

    def reset_conv_parameters(self, init_method="standard") -> None:
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                self.reset_parameters(m, init_method)

    def get_bop_params(self):
        bop_params = []
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                bop_params += list(m.parameters())
        return bop_params

    def get_non_bop_params(self):
        non_bop_params = []
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.BatchNorm2d,)):
                non_bop_params += list(m.parameters())
        return non_bop_params

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        if self.layer4 is not None:
            out = self.layer4(out)
        # out = F.avg_pool2d(out, 4)
        out = F.avg_pool2d(out, out.size()[3])
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


def PsgSeedResNet20(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(BasicBlock, [3,3,3], in_planes=16, num_classes=num_classes, init_method=init_method)

def PsgSeedResNet18(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(BasicBlock, [2,2,2,2], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet34(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(BasicBlock, [3,4,6,3], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet50(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(Bottleneck, [3,4,6,3], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet101(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(Bottleneck, [3,4,23,3], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet152(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(Bottleneck, [3,8,36,3], num_classes=num_classes, init_method=init_method)


def test():
    net = ResNet18()
    y = net(torch.randn(1,3,32,32))
    print(y.size())

# test()

[((971, 1487), 'models.masked_psg_seed_conv.PredictiveSeedConv2d', 'PredictiveSeedConv2d', (['in_planes', 'out_planes'], {'kernel_size': '(1)', 'stride': 'stride', 'padding': '(0)', 'bias': '(False)', 'num_bits': 'NUM_BITS', 'num_bits_weight': 'NUM_BITS_WEIGHT', 'num_bits_grad': 'NUM_BITS_GRAD', 'biprecision': 'BIPRECISION', 'input_signed': 'input_signed', 'predictive_forward': 'predictive_forward', 'predictive_backward': 'PREDICTIVE_BACKWARD', 'msb_bits': 'MSB_BITS', 'msb_bits_weight': 'MSB_BITS_WEIGHT', 'msb_bits_grad': 'MSB_BITS_GRAD', 'threshold': 'THRESHOLD', 'sparsify': 'SPARSIFY', 'sign': 'SIGN', 'writer': 'WRITER', 'writer_prefix': 'writer_prefix'}), '(in_planes, out_planes, kernel_size=1, stride=stride,\n    padding=0, bias=False, num_bits=NUM_BITS, num_bits_weight=\n    NUM_BITS_WEIGHT, num_bits_grad=NUM_BITS_GRAD, biprecision=BIPRECISION,\n    input_signed=input_signed, predictive_forward=predictive_forward,\n    predictive_backward=PREDICTIVE_BACKWARD, msb_bits=MSB_BITS,\n    msb_bits_weight=MSB_BITS_WEIGHT, msb_bits_grad=MSB_BITS_GRAD, threshold\n    =THRESHOLD, sparsify=SPARSIFY, sign=SIGN, writer=WRITER, writer_prefix=\n    writer_prefix)\n', (991, 1487), False, 'from models.masked_psg_seed_conv import PredictiveSeedConv2d\n'), ((1738, 2254), 'models.masked_psg_seed_conv.PredictiveSeedConv2d', 'PredictiveSeedConv2d', (['in_planes', 'out_planes'], {'kernel_size': '(3)', 'stride': 'stride', 'padding': '(1)', 'bias': '(False)', 'num_bits': 'NUM_BITS', 'num_bits_weight': 'NUM_BITS_WEIGHT', 'num_bits_grad': 'NUM_BITS_GRAD', 'biprecision': 'BIPRECISION', 'input_signed': 'input_signed', 'predictive_forward': 'predictive_forward', 'predictive_backward': 'PREDICTIVE_BACKWARD', 'msb_bits': 'MSB_BITS', 'msb_bits_weight': 'MSB_BITS_WEIGHT', 'msb_bits_grad': 'MSB_BITS_GRAD', 'threshold': 'THRESHOLD', 'sparsify': 'SPARSIFY', 'sign': 'SIGN', 'writer': 'WRITER', 'writer_prefix': 'writer_prefix'}), '(in_planes, out_planes, kernel_size=3, stride=stride,\n    padding=1, bias=False, num_bits=NUM_BITS, num_bits_weight=\n    NUM_BITS_WEIGHT, num_bits_grad=NUM_BITS_GRAD, biprecision=BIPRECISION,\n    input_signed=input_signed, predictive_forward=predictive_forward,\n    predictive_backward=PREDICTIVE_BACKWARD, msb_bits=MSB_BITS,\n    msb_bits_weight=MSB_BITS_WEIGHT, msb_bits_grad=MSB_BITS_GRAD, threshold\n    =THRESHOLD, sparsify=SPARSIFY, sign=SIGN, writer=WRITER, writer_prefix=\n    writer_prefix)\n', (1758, 2254), False, 'from models.masked_psg_seed_conv import PredictiveSeedConv2d\n'), ((2575, 2597), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['planes'], {}), '(planes)\n', (2589, 2597), True, 'import torch.nn as nn\n'), ((2738, 2760), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['planes'], {}), '(planes)\n', (2752, 2760), True, 'import torch.nn as nn\n'), ((2786, 2801), 'torch.nn.Sequential', 'nn.Sequential', ([], {}), '()\n', (2799, 2801), True, 'import torch.nn as nn\n'), ((3379, 3390), 'torch.nn.functional.relu', 'F.relu', (['out'], {}), '(out)\n', (3385, 3390), True, 'import torch.nn.functional as F\n'), ((3778, 3800), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['planes'], {}), '(planes)\n', (3792, 3800), True, 'import torch.nn as nn\n'), ((4048, 4070), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['planes'], {}), '(planes)\n', (4062, 4070), True, 'import torch.nn as nn\n'), ((4317, 4356), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(self.expansion * planes)'], {}), '(self.expansion * planes)\n', (4331, 4356), True, 'import torch.nn as nn\n'), ((4380, 4395), 'torch.nn.Sequential', 'nn.Sequential', ([], {}), '()\n', (4393, 4395), True, 'import torch.nn as nn\n'), ((5021, 5032), 'torch.nn.functional.relu', 'F.relu', (['out'], {}), '(out)\n', (5027, 5032), True, 'import torch.nn.functional as F\n'), ((5393, 5423), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['self.in_planes'], {}), '(self.in_planes)\n', (5407, 5423), True, 'import torch.nn as nn\n'), ((9475, 9497), 'torch.nn.Sequential', 'nn.Sequential', (['*layers'], {}), '(*layers)\n', (9488, 9497), True, 'import torch.nn as nn\n'), ((13663, 13688), 'torch.randn', 'torch.randn', (['(1)', '(3)', '(32)', '(32)'], {}), '(1, 3, 32, 32)\n', (13674, 13688), False, 'import torch\n'), ((5895, 5940), 'torch.nn.Linear', 'nn.Linear', (['(512 * block.expansion)', 'num_classes'], {}), '(512 * block.expansion, num_classes)\n', (5904, 5940), True, 'import torch.nn as nn\n'), ((6632, 6687), 'torch.nn.init._calculate_correct_fan', 'nn.init._calculate_correct_fan', (['module.weight', '"""fan_in"""'], {}), "(module.weight, 'fan_in')\n", (6662, 6687), True, 'import torch.nn as nn\n'), ((6707, 6737), 'torch.nn.init.calculate_gain', 'nn.init.calculate_gain', (['"""relu"""'], {}), "('relu')\n", (6729, 6737), True, 'import torch.nn as nn\n'), ((3168, 3207), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(self.expansion * planes)'], {}), '(self.expansion * planes)\n', (3182, 3207), True, 'import torch.nn as nn\n'), ((4762, 4801), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(self.expansion * planes)'], {}), '(self.expansion * planes)\n', (4776, 4801), True, 'import torch.nn as nn\n'), ((6343, 6369), 'torch.nn.Linear', 'nn.Linear', (['(64)', 'num_classes'], {}), '(64, num_classes)\n', (6352, 6369), True, 'import torch.nn as nn\n'), ((6763, 6777), 'math.sqrt', 'math.sqrt', (['fan'], {}), '(fan)\n', (6772, 6777), False, 'import math\n'), ((6795, 6810), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (6808, 6810), False, 'import torch\n'), ((6956, 7011), 'torch.nn.init._calculate_correct_fan', 'nn.init._calculate_correct_fan', (['module.weight', '"""fan_in"""'], {}), "(module.weight, 'fan_in')\n", (6986, 7011), True, 'import torch.nn as nn\n'), ((7031, 7061), 'torch.nn.init.calculate_gain', 'nn.init.calculate_gain', (['"""relu"""'], {}), "('relu')\n", (7053, 7061), True, 'import torch.nn as nn\n'), ((7087, 7101), 'math.sqrt', 'math.sqrt', (['fan'], {}), '(fan)\n', (7096, 7101), False, 'import math\n'), ((7119, 7134), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (7132, 7134), False, 'import torch\n'), ((7273, 7347), 'torch.nn.init.kaiming_normal_', 'nn.init.kaiming_normal_', (['module.weight'], {'mode': '"""fan_in"""', 'nonlinearity': '"""relu"""'}), "(module.weight, mode='fan_in', nonlinearity='relu')\n", (7296, 7347), True, 'import torch.nn as nn\n'), ((7173, 7208), 'torch.ones_like', 'torch.ones_like', (['module.weight.data'], {}), '(module.weight.data)\n', (7188, 7208), False, 'import torch\n'), ((7407, 7482), 'torch.nn.init.kaiming_uniform_', 'nn.init.kaiming_uniform_', (['module.weight'], {'mode': '"""fan_in"""', 'nonlinearity': '"""relu"""'}), "(module.weight, mode='fan_in', nonlinearity='relu')\n", (7431, 7482), True, 'import torch.nn as nn\n'), ((7548, 7603), 'torch.nn.init._calculate_correct_fan', 'nn.init._calculate_correct_fan', (['module.weight', '"""fan_in"""'], {}), "(module.weight, 'fan_in')\n", (7578, 7603), True, 'import torch.nn as nn\n'), ((7623, 7653), 'torch.nn.init.calculate_gain', 'nn.init.calculate_gain', (['"""relu"""'], {}), "('relu')\n", (7645, 7653), True, 'import torch.nn as nn\n'), ((7681, 7701), 'math.sqrt', 'math.sqrt', (['(2.0 * fan)'], {}), '(2.0 * fan)\n', (7690, 7701), False, 'import math\n'), ((7719, 7734), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (7732, 7734), False, 'import torch\n'), ((7765, 7830), 'numpy.random.laplace', 'np.random.laplace', ([], {'loc': '(0.0)', 'scale': 'scale', 'size': 'module.weight.shape'}), '(loc=0.0, scale=scale, size=module.weight.shape)\n', (7782, 7830), True, 'import numpy as np\n'), ((8002, 8039), 'torch.nn.init.xavier_normal_', 'nn.init.xavier_normal_', (['module.weight'], {}), '(module.weight)\n', (8024, 8039), True, 'import torch.nn as nn\n'), ((8117, 8169), 'torch.nn.init._calculate_fan_in_and_fan_out', 'nn.init._calculate_fan_in_and_fan_out', (['module.weight'], {}), '(module.weight)\n', (8154, 8169), True, 'import torch.nn as nn\n'), ((8246, 8261), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (8259, 8261), False, 'import torch\n'), ((8426, 8438), 'math.sqrt', 'math.sqrt', (['(5)'], {}), '(5)\n', (8435, 8438), False, 'import math\n'), ((7898, 7926), 'torch.from_numpy', 'torch.from_numpy', (['new_weight'], {}), '(new_weight)\n', (7914, 7926), False, 'import torch\n')]