text:
#  Copyright (c) 2020, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

import sys
import itertools
import numpy as np
from coremltools.models.utils import _python_version
from coremltools.models.utils import _macos_version
from coremltools.converters.mil import testing_reqs
from coremltools.converters.mil.testing_reqs import *
from .testing_utils import *
from coremltools import TensorType, ImageType, RangeDim

backends = testing_reqs.backends

torch = pytest.importorskip("torch")
pytestmark = pytest.mark.skipif(
    sys.version_info >= (3, 8), reason="Segfault with Python 3.8+"
)  # rdar://problem/65730375

torch.manual_seed(30)
np.random.seed(30)

class TestArgSort(TorchBaseTest):
    @pytest.mark.parametrize(
        "rank, axis, descending, backend",
        itertools.product(
            [rank for rank in range(1, 6)],
            [-1, 0],
            [True, False],
            backends
        )
    )
    def test_argsort(self, rank, axis, descending, backend):
        shape = tuple(np.random.randint(low=1, high=4, size=rank))
        model = ModuleWrapper(
            function=torch.argsort, kwargs={"dim": axis, "descending": descending}
        )
        TorchBaseTest.run_compare_torch(shape, model, backend=backend)

class TestBatchNorm(TorchBaseTest):
    @pytest.mark.parametrize(
        "num_features, eps, backend",
        itertools.product([5, 3, 2, 1], [0.1, 1e-05, 1e-9], backends),
    )
    def test_batchnorm(self, num_features, eps, backend):
        model = nn.BatchNorm2d(num_features, eps)
        self.run_compare_torch((6, num_features, 5, 5), model, backend=backend)

    @pytest.mark.parametrize("backend", backends)
    def test_batchnorm_1d(self, backend):
        class CRNNBase(nn.Module):
            def __init__(self, ch_in, ch_out, kernel_size=3, use_bn=True):
                super(CRNNBase, self).__init__()
                self.conv = nn.Conv1d(ch_in, ch_out, kernel_size=kernel_size)
                self.norm = nn.BatchNorm1d(ch_out)

            def forward(self, x):
                x = self.conv(x)
                x = self.norm(x)
                return x
        model = CRNNBase(ch_in=6, ch_out=16)
        self.run_compare_torch((1, 6, 15), model, backend=backend)

    @pytest.mark.parametrize(
        "shape, eps, backend",
        itertools.product([(4, 8), (1, 5), (5, 10), (6, 1)], [0.1, 1e-05, 1e-9], backends),
    )
    def test_batchnorm1d_rank2(self, shape, eps, backend):
        N,C = shape
        batchnorm = nn.BatchNorm1d(C, eps=eps).eval()
        self.run_compare_torch(
            (N, C), batchnorm, backend=backend,
        )

    @pytest.mark.parametrize(
        "shape, eps, backend",
        itertools.product([(4, 8, 2), (1, 5, 3), (5, 10, 1), (6, 1, 4)], [0.1, 1e-05, 1e-9], backends),
    )
    def test_batchnorm1d_rank3(self, shape, eps, backend):
        N,C,L = shape
        batchnorm = nn.BatchNorm1d(C, eps=eps).eval()
        self.run_compare_torch(
            (N, C, L), batchnorm, backend=backend,
        )

class TestInstanceNorm(TorchBaseTest):
    @pytest.mark.parametrize(
        "num_features, eps, backend",
        itertools.product([5, 3, 2, 1], [0.1, 1e-05, 1e-09], backends),
    )
    def test_instancenorm(self, num_features, eps, backend):
        if backend == "nn_proto" and eps == 1e-09:
            return
        model = nn.InstanceNorm2d(num_features, eps)
        self.run_compare_torch((6, num_features, 5, 5), model, backend=backend)

    @pytest.mark.parametrize("num_features, backend",
                             itertools.product([5, 3, 2, 1], backends),
    )
    def test_instancenorm_1d(self, num_features, backend):
        model = nn.InstanceNorm1d(num_features)
        self.run_compare_torch((6, num_features, 10), model, backend=backend)

class TestGroupNorm(TorchBaseTest):
    @pytest.mark.parametrize(
        "group_features, eps,affine, backend",
        itertools.product([(16,32), (32,64), (1,1)], [0.1, 1e-05, 1e-09],[True, False], backends),
    )
    def test_groupnorm(self, group_features, eps, affine, backend):
        if backend == "nn_proto" and eps == 1e-09:
            return
        model = nn.GroupNorm(group_features[0],group_features[1], eps=eps, affine=affine)
        self.run_compare_torch((6, group_features[1], 5, 5), model, backend=backend)

class TestLinear(TorchBaseTest):
    @pytest.mark.parametrize(
        "in_features, out_features, backend",
        itertools.product([10, 25, 100], [3, 6], backends),
    )
    def test_addmm(self, in_features, out_features, backend):
        model = nn.Linear(in_features, out_features)
        self.run_compare_torch((1, in_features), model, backend=backend)

    @pytest.mark.parametrize(
        "in_features, out_features, backend",
        itertools.product([5], [10], backends),
    )
    def test_linear_rank1_input(self, in_features, out_features, backend):
        model = nn.Linear(in_features, out_features)
        self.run_compare_torch((in_features,), model, backend=backend)

class TestConv(TorchBaseTest):
    @pytest.mark.parametrize(
        "height, width, in_channels, out_channels, kernel_size, stride, padding, dilation, backend",
        itertools.product(
            [5, 6], [5, 7], [1, 3], [1, 3], [1, 3], [1, 3], [1, 3], [1, 3], backends
        ),
    )
    def test_convolution2d(
        self,
        height,
        width,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        dilation,
        backend,
        groups=1,
    ):
        model = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
        )
        self.run_compare_torch((1, in_channels, height, width), model,
                           backend=backend)

class TestConvTranspose(TorchBaseTest):
    @pytest.mark.parametrize(
        "width, in_channels, out_channels, kernel_size, stride, padding, dilation, backend",
        itertools.product(
            [5, 7], [1, 3], [1, 3], [1, 3], [2, 3], [0, 1], [1, 3], backends
        ),
    )
    def test_convolution_transpose1d(
            self,
            width,
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            dilation,
            backend,
            groups=1,
    ):
        model = nn.ConvTranspose1d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups
        )
        self.run_compare_torch((1, in_channels, width), model, backend=backend)

    @pytest.mark.parametrize(
        "height, width, in_channels, out_channels, kernel_size, stride, padding, dilation, backend",
        itertools.product(
            [5, 6], [5, 7], [1, 3], [1, 3], [1, 3], [2, 3], [0, 1], [1, 3], backends
        ),
    )
    def test_convolution_transpose2d(
        self,
        height,
        width,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        dilation,
        backend,
        groups=1,
    ):
        model = nn.ConvTranspose2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
        )
        self.run_compare_torch((1, in_channels, height, width), model,
                           backend=backend)

    # TODO: rdar://65588783 ([PyTorch] Define and error out on unsupported configuration for output_padding)
    # TODO: rdar://65550420 (Add Image Resizing (crop, upsample, resize_bilinear) layers to the MIL backend)
    @pytest.mark.parametrize(
        "height, width, in_channels, out_channels, kernel_size, stride, padding, dilation, output_padding, backend",
        list(
            itertools.product(
                [10],
                [10],
                [1, 3],
                [1, 3],
                [1, 3],
                [1, 2, 3],
                [1, 3],
                [1, 2],
                [1, 2, (1, 2)],
                ["nn_proto"],
            )
        )
        + [
            pytest.param(
                5, 5, 1, 1, 3, 4, 1, 1, 2, "nn_proto", marks=pytest.mark.xfail
            ),
            pytest.param(
                5, 5, 1, 1, 3, 2, 1, 3, 2, "nn_proto", marks=pytest.mark.xfail
            ),
        ],
    )
    def test_convolution_transpose2d_output_padding(
        self,
        height,
        width,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        dilation,
        output_padding,
        backend,
        groups=1,
    ):

        # Output padding must be less than either stride or dilation
        # Skip testing invalid combinations
        if isinstance(output_padding, int):
            if output_padding >= stride and output_padding >= dilation:
                return
        elif isinstance(output_padding, tuple):
            for _output_padding in output_padding:
                if _output_padding >= stride and _output_padding >= dilation:
                    return

        model = nn.ConvTranspose2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            output_padding=output_padding,
        )
        self.run_compare_torch((1, in_channels, height, width), model,
                           backend=backend)

    @pytest.mark.parametrize(
        "depth, height, width, in_channels, out_channels, kernel_size, stride, padding, dilation, backend",
        itertools.product(
            [3, 4], [5, 6], [5, 7], [1, 3], [1, 3], [1, 3], [2, 3], [0, 1], [1, 3], backends
        ),
    )
    @pytest.mark.skip(reason="rdar://65198011 (Re-enable Conv3dTranspose and DynamicTile unit tests)")
    def test_convolution_transpose3d(
             self,
             depth,
             height,
             width,
             in_channels,
             out_channels,
             kernel_size,
             stride,
             padding,
             dilation,
             backend,
    ):
        model = nn.ConvTranspose3d(
             in_channels=in_channels,
             out_channels=out_channels,
             kernel_size=kernel_size,
             stride=stride,
             padding=padding,
             dilation=dilation,
         )
        self.run_compare_torch((1, in_channels, depth, height, width), model,
                           backend=backend)

class TestCond(TorchBaseTest):
    @pytest.mark.parametrize("backend", backends)
    def test_cond(self, backend):
        in_features = 1
        out_features = 2
        class TestNet(nn.Module):
            def forward(self, x):
                if torch.squeeze(x) < 10.:
                    return x*10.
                else:
                    return x*2.

        model = TestNet().eval()
        torch_model = torch.jit.script(model)

        self.run_compare_torch(torch.tensor([1.]), torch_model,
            input_as_shape=False, backend=backend)
        self.run_compare_torch(torch.tensor([11.]), torch_model,
            input_as_shape=False, backend=backend)

class TestLoop(TorchBaseTest):
    @pytest.mark.parametrize("backend", backends)
    def test_for_loop(self, backend):
        class TestLayer(nn.Module):
            def __init__(self):
                super(TestLayer, self).__init__()

            def forward(self, x):
                x = 2.0 * x
                return x

        class TestNet(nn.Module):
            input_size = (64,)

            def __init__(self):
                super(TestNet, self).__init__()
                layer = TestLayer()
                self.layer = torch.jit.trace(layer, torch.rand(self.input_size))

            def forward(self, x):
                for _ in range(7):
                    x = self.layer(x)
                return x

        model = TestNet().eval()
        torch_model = torch.jit.script(model)

        self.run_compare_torch(model.input_size, torch_model, backend=backend)

    @pytest.mark.parametrize("backend", backends)
    def test_while_loop(self, backend):
        class TestLayer(nn.Module):
            def __init__(self):
                super(TestLayer, self).__init__()

            def forward(self, x):
                x = 0.5 * x
                return x

        class TestNet(nn.Module):
            input_size = (1,)

            def __init__(self):
                super(TestNet, self).__init__()
                layer = TestLayer()
                self.layer = torch.jit.trace(layer, torch.rand(self.input_size))

            def forward(self, x):
                while x > 0.01:
                    x = self.layer(x)
                return x

        model = TestNet().eval()
        torch_model = torch.jit.script(model)

        self.run_compare_torch(model.input_size, torch_model, backend=backend)

class TestUpsample(TorchBaseTest):
    @pytest.mark.parametrize(
        "output_size, align_corners, backend",
        [
            x
            for x in itertools.product(
                [(10, 10), (1, 1), (20, 20), (2, 3), (190, 170)],
                [True, False],
                backends,
            )
        ],
    )
    def test_upsample_bilinear2d_with_output_size(
        self, output_size, align_corners, backend
    ):
        input_shape = (1, 3, 10, 10)
        model = ModuleWrapper(
            nn.functional.interpolate,
            {"size": output_size, "mode": "bilinear", "align_corners": align_corners,},
        )
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "scales_h, scales_w, align_corners, backend",
        [
            x
            for x in itertools.product(
                [2, 3, 4.5], [4, 5, 5.5], [True, False], backends
            )
        ],
    )
    def test_upsample_bilinear2d_with_scales(
        self, scales_h, scales_w, align_corners, backend
    ):
        input_shape = (1, 3, 10, 10)
        model = ModuleWrapper(
            nn.functional.interpolate,
            {
                "scale_factor": (scales_h, scales_w),
                "mode": "bilinear",
                "align_corners": align_corners,
            },
        )
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "output_size, backend",
        [
            x
            for x in itertools.product(
                [(10, 10), (30, 20), (20, 20), (20, 30), (190, 170)], backends
            )
        ],
    )
    def test_upsample_nearest2d_with_output_size(self, output_size, backend):
        input_shape = (1, 3, 10, 10)
        model = ModuleWrapper(
            nn.functional.interpolate, {"size": output_size, "mode": "nearest"},
        )
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "scales_h, scales_w, backend",
        [x for x in itertools.product([2, 3, 5, 4.5], [4, 5, 2, 5.5], backends)],
    )
    def test_upsample_nearest2d_with_scales(self, scales_h, scales_w, backend):
        if backend == "nn_proto":
            if isinstance(scales_h, float) or isinstance(scales_w, float):
                return  # Skip fractional scale factors tests for nn_proto

        input_shape = (1, 3, 10, 10)
        model = ModuleWrapper(
            nn.functional.interpolate,
            {"scale_factor": (scales_h, scales_w), "mode": "nearest"},
        )
        self.run_compare_torch(input_shape, model, backend=backend)

class TestBranch(TorchBaseTest):
    @pytest.mark.parametrize("backend", backends)
    def test_if(self, backend):
        class TestLayer(nn.Module):
            def __init__(self):
                super(TestLayer, self).__init__()

            def forward(self, x):
                x = torch.mean(x)
                return x

        class TestNet(nn.Module):
            input_size = (64,)

            def __init__(self):
                super(TestNet, self).__init__()
                layer = TestLayer()
                self.layer = torch.jit.trace(layer, torch.rand(self.input_size))

            def forward(self, x):
                m = self.layer(x)
                if m < 0:
                    scale = -2.0
                else:
                    scale = 2.0
                x = scale * x
                return x

        model = TestNet().eval()
        torch_model = torch.jit.script(model)

        self.run_compare_torch(model.input_size, torch_model, backend=backend)

class TestAvgPool(TorchBaseTest):
    # rdar://66066001 (PyTorch converter: enable ceil_mode=True tests for pooling ops)

    @pytest.mark.parametrize(
        "input_shape, kernel_size, stride, padding, ceil_mode, include_pad, backend",
        itertools.product(
            [(1, 3, 15), (1, 1, 7), (1, 3, 10)],
            [1, 2, 3],
            [1, 2],
            [0, 1],
            [False],
            [True, False],
            backends,
        ),
    )
    def test_avg_pool1d(
        self, input_shape, kernel_size, stride, padding, ceil_mode, include_pad, backend
    ):
        if padding > kernel_size / 2:
            return
        model = nn.AvgPool1d(
            kernel_size,
            stride,
            padding,
            ceil_mode=ceil_mode,
            count_include_pad=include_pad,
        )
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "input_shape, kernel_size, stride, padding, ceil_mode, include_pad, backend",
        itertools.product(
            [(1, 3, 15, 15), (1, 1, 7, 7), (1, 3, 10, 10)],
            [1, 2, 3],
            [1, 2],
            [0, 1],
            [False],
            [True, False],
            backends,
        ),
    )
    def test_avg_pool2d(
        self, input_shape, kernel_size, stride, padding, ceil_mode, include_pad, backend
    ):
        if padding > kernel_size / 2:
            return
        model = nn.AvgPool2d(
            kernel_size,
            stride,
            padding,
            ceil_mode=ceil_mode,
            count_include_pad=include_pad,
        )
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "input_shape, kernel_size, stride, padding, ceil_mode, include_pad, backend",
        itertools.product(
            [(1, 3, 11, 3, 11), (1, 1, 7, 4, 7), (1, 3, 6, 6, 3)],
            [1, 2, 3],
            [1, 2],
            [0, 1],
            [False],
            [True, False],
            backends,
        ),
    )
    def test_avg_pool3d(
        self, input_shape, kernel_size, stride, padding, ceil_mode, include_pad, backend
    ):
        if padding > kernel_size / 2:
            return
        model = nn.AvgPool3d(
            kernel_size,
            stride,
            padding,
            ceil_mode=ceil_mode,
            count_include_pad=include_pad,
        )
        self.run_compare_torch(input_shape, model, backend=backend)

class TestAdaptiveMaxPool(TorchBaseTest):
    @pytest.mark.parametrize(
        "output_size, magnification, delta, depth, backend",
        itertools.product(
            [(1,1), (3,2),(3,6),(32,32)],
            [1,2,4,5,6,7],
            [0,11],
            [1,2,3],
            backends,
        ),
    )
    def test_adaptive_max_pool2d(
            self, output_size, magnification, delta, depth, backend
    ):
        # input_size = output_size * magnification + delta
        input_size = (delta + magnification * output_size[0], delta + magnification * output_size[1])
        # since coremltools reproduces PyTorch's kernel sizes and
        # offsets for adaptive pooling layers only when input_size is
        # a multiple of output_size, we expect failures otherwise
        if not (input_size[0] % output_size[0]  == 0 and input_size[1] % output_size[1] == 0):
            pytest.xfail("Test should fail because input_size is not a multiple of output_size")
        n = 1
        in_shape = (n,depth) + input_size
        model = nn.AdaptiveMaxPool2d(
            output_size
        )
        self.run_compare_torch(in_shape, model, backend=backend)

class TestMaxPool(TorchBaseTest):
    # rdar://66066001 (PyTorch converter: enable ceil_mode=True tests for pooling ops)

    @pytest.mark.parametrize(
        "input_shape, kernel_size, stride, padding, ceil_mode, backend",
        itertools.product(
            [(1, 3, 15), (1, 1, 7), (1, 3, 10)],
            [1, 2, 3],
            [1, 2],
            [0, 1],
            [False],
            backends,
        ),
    )
    def test_max_pool1d(
        self, input_shape, kernel_size, stride, padding, ceil_mode, backend
    ):
        if padding > kernel_size / 2:
            return
        model = nn.MaxPool1d(
            kernel_size,
            stride,
            padding,
            dilation=1,
            return_indices=False,
            ceil_mode=ceil_mode,
        )
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "input_shape, kernel_size, stride, padding, ceil_mode, backend",
        itertools.product(
            [(1, 3, 15, 15), (1, 1, 7, 7), (1, 3, 10, 10)],
            [1, 2, 3],
            [1, 2],
            [0, 1],
            [False],
            backends,
        ),
    )
    def test_max_pool2d(
        self, input_shape, kernel_size, stride, padding, ceil_mode, backend
    ):
        if padding > kernel_size / 2:
            return
        model = nn.MaxPool2d(
            kernel_size,
            stride,
            padding,
            dilation=1,
            return_indices=False,
            ceil_mode=ceil_mode,
        )
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "input_shape, kernel_size, stride, padding, ceil_mode, backend",
        itertools.product(
            [(1, 3, 11, 3, 11), (1, 1, 7, 4, 7), (1, 3, 6, 6, 3)],
            [1, 2, 3],
            [1, 2],
            [0, 1],
            [False],
            backends,
        ),
    )
    def test_max_pool3d(
        self, input_shape, kernel_size, stride, padding, ceil_mode, backend
    ):
        if padding > kernel_size / 2:
            return
        model = nn.MaxPool3d(
            kernel_size,
            stride,
            padding,
            dilation=1,
            return_indices=False,
            ceil_mode=ceil_mode,
        )
        self.run_compare_torch(input_shape, model, backend=backend)

class TestLSTM(TorchBaseTest):
    @pytest.mark.parametrize(
        "input_size, hidden_size, num_layers, bias, batch_first, dropout, bidirectional, backend",
        itertools.product(
            [7], [5], [1, 2, 3], [True, False], [False, True], [0.3], [False, True], backends
        ),
    )
    def test_lstm(
        self,
        input_size,
        hidden_size,
        num_layers,
        bias,
        batch_first,
        dropout,
        bidirectional,
        backend,
    ):
        model = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bias=bias,
            batch_first=batch_first,
            dropout=dropout,
            bidirectional=bidirectional,
        )
        SEQUENCE_LENGTH = 3
        BATCH_SIZE = 2
        model.eval()

        num_directions = int(bidirectional) + 1

        if batch_first:
            _input = torch.randn(BATCH_SIZE, SEQUENCE_LENGTH, input_size)
        else:
            _input = torch.randn(SEQUENCE_LENGTH, BATCH_SIZE, input_size)

        h0 = torch.randn(num_layers * num_directions, BATCH_SIZE, hidden_size)
        c0 = torch.randn(num_layers * num_directions, BATCH_SIZE, hidden_size)

        inputs = (_input, (h0, c0))
        expected_results = model(*inputs)
        self.run_compare_torch(
            inputs, model, expected_results, input_as_shape=False, backend=backend,
        )

# Workaround for GitHub Issue #824
# i.e. the return h_n/c_n for a converted BLSTM are mangled.
# Therefore, just look at output 'y' (for now) which is correct.
class StripCellAndHidden(nn.Module):
    def __init__(self,flagReturnTuple_):
        super(StripCellAndHidden, self).__init__()
        self.flagReturnTuple = flagReturnTuple_

    def forward(self,x):
        # Pass tuple, not tensor, to avoid issue in coremltools/converters/mil/frontend/torch/test/testing_utils.py on "if not expected_results:"
        # Pass tensor when we need input for LSTM #2 as part of nn.Sequential()
        return tuple(x[0]) if self.flagReturnTuple else x[0]

# Check GitHub Issue #810, assume num_layers == 2 and bidirectional == True
class TestStackedBLSTM(TorchBaseTest):
    @pytest.mark.parametrize(
        "input_size, hidden_size, num_layers, bias, batch_first, dropout, bidirectional, backend",
        itertools.product([7], [5], [2], [True, False], [True, False], [0.3], [True], backends),
    )
    def test_lstm(
        self,
        input_size,
        hidden_size,
        num_layers,
        bias,
        batch_first,
        dropout,
        bidirectional,
        backend,
    ):
        model = nn.Sequential(
            nn.LSTM(
                input_size=input_size,
                hidden_size=hidden_size,
                num_layers=1,
                bias=bias,
                batch_first=batch_first,
                dropout=dropout,
                bidirectional=True),
            StripCellAndHidden(False),
            nn.LSTM(
                input_size=2*hidden_size,
                hidden_size=hidden_size,
                num_layers=1,
                bias=bias,
                batch_first=batch_first,
                dropout=dropout,
                bidirectional=True),
            StripCellAndHidden(True)
        )

        SEQUENCE_LENGTH = 3
        BATCH_SIZE = 2

        num_directions = int(bidirectional) + 1

        # (seq_len, batch, input_size)
        if batch_first:
            _input = torch.rand(BATCH_SIZE, SEQUENCE_LENGTH, input_size)
        else:
            _input = torch.randn(SEQUENCE_LENGTH, BATCH_SIZE, input_size)

        # Do not use h_0/c_0 input and do not check h_n/c_n output, GitHub Issue #824
        expected_results = model(_input)

        self.run_compare_torch(_input, model, expected_results,
                           input_as_shape=False, backend=backend)

class TestConcat(TorchBaseTest):
    # This tests an edge case where the list of tensors to concatenate only
    # has one item. NN throws an error for this case, hence why we have to
    # run through the full conversion process to test it.
    @pytest.mark.parametrize("backend", backends)
    def test_cat(self, backend):
        class TestNet(nn.Module):
            def __init__(self):
                super(TestNet, self).__init__()

            def forward(self, x):
                x = torch.cat((x,), axis=1)
                return x

        model = TestNet()
        self.run_compare_torch((1, 3, 16, 16), model, backend=backend)

class TestReduction(TorchBaseTest):
    @pytest.mark.parametrize(
        "input_shape, dim, keepdim, backend",
        itertools.product([(2, 2), (1, 1)], [0, 1], [True, False], backends),
    )
    def test_max(self, input_shape, dim, keepdim, backend):
        class TestMax(nn.Module):
            def __init__(self):
                super(TestMax, self).__init__()

            def forward(self, x):
                return torch.max(x, dim=dim, keepdim=keepdim)

        input_data = torch.rand(input_shape)
        model = TestMax()
        # TODO: Expected results are flipped due to naming issue:
        # rdar://62681982 (Determine the output names of MLModels)
        expected_results = model(input_data)[::-1]
        self.run_compare_torch(
            input_data,
            model,
            expected_results=expected_results,
            input_as_shape=False,
            backend=backend,
        )

class TestLayerNorm(TorchBaseTest):
    @pytest.mark.parametrize(
        "input_shape, eps, backend",
        itertools.product([(1, 3, 15, 15), (1, 1, 1, 1)], [1e-5, 1e-9], backends),
    )
    def test_layer_norm(self, input_shape, eps, backend):
        model = nn.LayerNorm(input_shape, eps=eps)
        self.run_compare_torch(input_shape, model, backend=backend)

class TestPixelShuffle(TorchBaseTest):
    @pytest.mark.parametrize(
        "batch_size, CHW, r, backend",
        itertools.product([1, 3], [(1, 4, 4), (3, 2, 3)], [2, 4], backends),
    )
    def test_pixel_shuffle(self, batch_size, CHW, r, backend):
        C, H, W = CHW
        input_shape = (batch_size, C * r * r, H, W)
        model = nn.PixelShuffle(upscale_factor=r)
        self.run_compare_torch(input_shape, model, backend=backend)

class TestExpand(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, shapes",
        itertools.product(
            backends,
            [[(2, 1), (2, 2)], [(3, 1), (-1, 4)], [(1, 3, 4, 4), (3, 3, 4, 4)]]
        ),
    )
    def test_expand(self, backend, shapes):
        input_shape, output_shape = shapes

        class TestModel(torch.nn.Module):
            def forward(self, x):
                return x.expand(*output_shape)

        model = TestModel()

        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "backend, input_shapes",
        itertools.product(
            backends,
            [[(2, 1), (2, 2)], [(3, 1), (3, 4)], [(1, 3, 4, 4), (3, 3, 4, 4)]]
        ),
    )
    def test_expand_as(self, backend, input_shapes):
        class TestModel(torch.nn.Module):
            def forward(self, x, y):
                return x.expand_as(y)

        model = TestModel()

        self.run_compare_torch(input_shapes, model, backend=backend)

class TestExpandDims(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, rank_and_axis",
        itertools.product(
            backends,
            [
                (rank, axis)
                for rank in range(1, 5)
                for axis in range(-rank - 1, rank + 1)
            ],
        ),
    )
    def test_unsqueeze(self, backend, rank_and_axis):
        rank, axis = rank_and_axis
        input_shape = tuple(np.random.randint(low=2, high=10, size=rank))
        model = ModuleWrapper(function=torch.unsqueeze, kwargs={"dim": axis})
        self.run_compare_torch(input_shape, model, backend=backend)

class TestSqueeze(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, rank_and_axis",
        itertools.product(
            backends,
            [(2, 1), (2, 0), (3, 1), (3, None), (4, None), (4, 2), (5, None), (5, -1),],
        ),
    )
    def test_squeeze(self, backend, rank_and_axis):
        rank, axis = rank_and_axis
        input_shape = list(np.random.randint(low=2, high=10, size=rank))
        if axis is not None:
            input_shape[axis] = 1
        else:
            input_shape[0] = 1
        input_shape = tuple(input_shape)
        model = ModuleWrapper(
            function=torch.squeeze, kwargs={"dim": axis} if axis else {}
        )
        self.run_compare_torch(input_shape, model, backend=backend)

class TestCumSum(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, axis",
        itertools.product(
            backends,
            [-1, 0, 1, 2, 3],
        ),
    )
    def test_cumsum(self, backend, axis):
        input_shape = list(np.random.randint(low=2, high=10, size=4))
        input_shape = tuple(input_shape)
        model = ModuleWrapper(
            function=torch.cumsum, kwargs={"dim": axis}
        )
        self.run_compare_torch(input_shape, model, backend=backend)

class TestReshape(TorchBaseTest):
    # TODO: <rdar://66239973> Add dynamic & rank preserving reshape tests for pytorch
    @pytest.mark.parametrize(
        "backend, output_shape",
        itertools.product(backends, [(3, 2), (2, -1), (2, 1, 1, 3),],),
    )
    def test_reshape(self, backend, output_shape):
        input_shape = (2, 3)
        model = ModuleWrapper(function=torch.reshape, kwargs={"shape": output_shape})
        self.run_compare_torch(input_shape, model, backend=backend)

class TestFlatten(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, start_dim",
        itertools.product(backends, [2,-2],),
    )
    def test_reshape(self, backend, start_dim):
        input_shape = (2, 3, 4, 5)
        model = ModuleWrapper(function=torch.flatten, kwargs={"start_dim": start_dim})
        self.run_compare_torch(input_shape, model, backend=backend)

class TestGather(TorchBaseTest):
    @pytest.mark.xfail(
        reason="Load constant not copied properly for integer valued constants. Enable after eng/PR-65551506 is merged",
        run=False,
    )
    @pytest.mark.parametrize(
        "rank_and_axis, backend",
        itertools.product([(i, j) for i in range(1, 6) for j in range(0, i)], backends),
    )
    def test_gather_along_axis(self, rank_and_axis, backend):
        rank, axis = rank_and_axis
        params_shape = np.random.randint(low=2, high=5, size=rank)
        indices_shape = np.copy(params_shape)
        indices_shape[axis] = np.random.randint(low=1, high=8)
        indices = np.random.randint(0, params_shape[axis], size=indices_shape)
        params_shape, indices_shape = tuple(params_shape), tuple(indices_shape)
        model = ModuleWrapper(
            function=torch.gather,
            kwargs={"dim": axis, "index": torch.from_numpy(indices)},
        )
        self.run_compare_torch([params_shape], model, backend=backend)

class TestActivation(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, rank", itertools.product(backends, range(1, 6)),
    )
    def test_relu(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.ReLU().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

        model = ModuleWrapper(nn.functional.relu_)
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank", itertools.product(backends, range(1, 6)),
    )
    def test_relu6(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.ReLU6().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, alpha", itertools.product(backends, [0.1, 0.25, 2.0]),
    )
    def test_prelu(self, backend, alpha):
        input_shape = tuple(np.random.randint(low=5, high=10, size=4))
        C = input_shape[1]
        model = nn.PReLU(C, alpha).eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank, alpha",
        itertools.product(backends, range(1, 6), [0.1, 2.0, 1.5]),
    )
    def test_leaky_relu(self, backend, rank, alpha):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.LeakyReLU(negative_slope=alpha).eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

        model = ModuleWrapper(nn.functional.leaky_relu_, {'negative_slope': alpha})
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank", itertools.product(backends, range(1, 6)),
    )
    def test_softmax(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.Softmax().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank, range_val",
        itertools.product(
            backends, range(1, 6), [(-1.0, 1.0), (0.0, 0.1), (1.0, 3.0), (-1.0, 6.0)]
        ),
    )
    def test_hardtanh(self, backend, rank, range_val):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.Hardtanh(range_val[0], range_val[1]).eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

        model = ModuleWrapper(nn.functional.hardtanh_,
                     {'min_val': range_val[0], 'max_val': range_val[1]})
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank, alpha",
        itertools.product(backends, range(1, 6), [0.1, 2.0, 1.5]),
    )
    def test_elu(self, backend, rank, alpha):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.ELU(alpha).eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    # rdar://problem/66557565
    @pytest.mark.parametrize(
        "backend, rank", itertools.product(['nn_proto'], range(1, 6)),
    )
    def test_gelu(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.GELU().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.skipif(_python_version() < (3, 6), reason="requires python 3.6")
    @pytest.mark.parametrize(
        "backend, rank", itertools.product(backends, range(1, 6)),
    )
    def test_erf(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))

        class ERFActivation(nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x):
                return torch.erf(x)

        model = ERFActivation().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank", itertools.product(backends, range(1, 6)),
    )
    def test_sigmoid(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.Sigmoid().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.skipif(_python_version() < (3, 6), reason="requires python 3.6")
    @pytest.mark.parametrize(
        "backend, rank", itertools.product(backends, range(1, 6)),
    )
    def test_sigmoid_hard(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.Hardsigmoid().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, beta, threshold", itertools.product(backends, [1, 2, 5], [5, 10, 20]),
    )
    @pytest.mark.skipif(
        _macos_version() <= (11,),
        reason="Parametric SoftPlus segfaults on macOS 10.15 and below. (rdar://problem/66555235)",
    )
    def test_softplus(self, backend, beta, threshold):
        input_shape = (1, 10, 5, 15)
        model = nn.Softplus(beta, threshold).eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    # rdar://problem/66557565
    @pytest.mark.parametrize(
        "backend, rank", itertools.product(['nn_proto'], range(1, 6)),
    )
    def test_softsign(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = nn.Softsign().eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

class TestElementWiseUnary(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, rank, op_string",
        itertools.product(
            backends,
            [4],
            [
                "abs",
                "acos",
                "asin",
                "atan",
                "ceil",
                "cos",
                "cosh",
                "exp",
                "floor",
                "round",
                "sin",
                "sinh",
                "sqrt",
                "square",
                "tan",
                "tanh",
                "sign",
            ],
        ),
    )
    def test_elementwise_no_params(self, backend, rank, op_string):
        if not contains_op(torch, op_string):
            return
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        op_func = getattr(torch, op_string)
        model = ModuleWrapper(function=op_func)
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    ## TODO (rdar://66577921): Needs to move to test_elementwise_no_params after backend is added
    @pytest.mark.parametrize(
        "backend, rank",
        itertools.product(
            ['nn_proto'],
            [4],
        ),
    )
    @pytest.mark.skipif(sys.version_info < (3, 6), reason="requires python3.6 or higher")
    def test_square(self, backend, rank):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = ModuleWrapper(function=torch.square)
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank, clamp_range",
        itertools.product(
            backends,
            [4],
            [(0.0, 1.0), (-1.0, 0.5), (0.2, 0.7), (None, 4.0), (-3.0, None)],
        ),
    )
    def test_clamp(self, backend, rank, clamp_range):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        params_dict = {}
        if clamp_range[0] is not None:
            params_dict["min"] = clamp_range[0]
        if clamp_range[1] is not None:
            params_dict["max"] = clamp_range[1]

        model = ModuleWrapper(torch.clamp, params_dict)
        self.run_compare_torch(
            input_shape, model, backend=backend, rand_range=(-5, 5)
        )

    @pytest.mark.parametrize(
        "backend, rank, threshold",
        itertools.product(
            ['nn_proto'], # rdar://66597974 Renable for all backends due to missing cast
            [4],
            [(0.0, 0.0), (0.5, 0.5), (0.5, 10), (0.9, 0.0)]
        ),
    )
    def test_threshold(self, backend, rank, threshold):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = torch.nn.Threshold(threshold[0], threshold[1]).eval()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )

    @pytest.mark.parametrize(
        "backend, rank, op_string",
        itertools.product(
            backends,
            [4],
            [
                "log",
                "rsqrt",
                "reciprocal",
            ],
        ),
    )
    def test_elementwise_numerically_stable(self, backend, rank, op_string):
        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        op_func = getattr(torch, op_string)
        model = ModuleWrapper(function=op_func)
        self.run_compare_torch(
            input_shape, model, backend=backend, rand_range=(20, 100)
        )

class TestMatMul(TorchBaseTest):

    @pytest.mark.parametrize("backend", backends)
    def test_bmm(self, backend):
        shape_x, shape_y = (3,4,5), (3,5,6)
        model = ModuleWrapper(function=torch.bmm)
        self.run_compare_torch(
            [shape_x, shape_y], model, backend=backend,
        )

class TestSplit(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, split_size_or_sections, dim",
        itertools.product(backends, [1, 2, [1, 4]], [0, -2]),
    )
    def test_split(self, backend, split_size_or_sections, dim):
        input_shape = (5, 2)
        model = ModuleWrapper(function=torch.split,
                              kwargs={"split_size_or_sections": split_size_or_sections, "dim": dim})
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize(
        "backend, split_sizes, dim",
        itertools.product(backends, [[1, 4], [3, 2]], [-1, -2]),
    )
    def test_split_with_sizes(self, backend, split_sizes, dim):
        input_shape = (5, 5)
        model = ModuleWrapper(function=torch.split_with_sizes,
                              kwargs={"split_sizes": split_sizes, "dim": dim})
        self.run_compare_torch(input_shape, model, backend=backend)

class TestUnbind(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, dim",
        itertools.product(backends,[0,1,2]),
    )
    def test_unbind(self, backend, dim):
        input_shape = (3, 3, 4)
        model = ModuleWrapper(function=torch.unbind,
                              kwargs={"dim": dim})
        self.run_compare_torch(input_shape, model, backend=backend)

class TestTranspose(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, rank, dims",
        itertools.product(backends, list(range(2, 6)),
                          [(0, 1), (-2, -1), (1, 0), (-1, -2)]),
    )
    def test(self, backend, rank, dims):
        input_shape = tuple(np.random.randint(low=1, high=4, size=rank))
        model = ModuleWrapper(function=torch.transpose,
                              kwargs={"dim0": dims[0], "dim1": dims[1]})
        self.run_compare_torch(input_shape, model, backend=backend)

class TestConstantPad(TorchBaseTest):
    @pytest.mark.parametrize(
        "input_shape, pad, backend",
        itertools.product([2,2], backends)
    )
    def test_constant_pad_nd(self, input_shape, backend):
        input_data = torch.rand(input_shape)
        model = torch.constant_pad_nd(input_data, (1,1,1,1))
        self.run_compare_torch(input_shape,model,backend=backend)

    @pytest.mark.parametrize("backend", backends)
    def test_constant_pad_1d(self, backend):
        input_shape = (3, 4, 5)
        model = torch.nn.ConstantPad1d((5, 6), 3.5).eval()
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize("backend", backends)
    def test_constant_pad_2d(self, backend):
        input_shape = (3, 4, 5, 6)
        model = torch.nn.ConstantPad2d((5, 6, 3, 8), 3.5).eval()
        self.run_compare_torch(input_shape, model, backend=backend)

    @pytest.mark.parametrize("backend", backends)
    def test_constant_pad_3d(self, backend):
        input_shape = (3, 4, 5, 6, 2)
        model = torch.nn.ConstantPad3d((5, 6, 3, 8, 2, 4), 3.5).eval()
        self.run_compare_torch(input_shape, model, backend=backend)

class TestTo(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend", backends,
    )
    def test_cast_bug(self, backend):
        class TestModel(torch.nn.Module):
            def forward(self, spans, embedding):
                spans = spans.float().relu().int()

                max1, _ = torch.max(spans, dim=1, keepdim=False)
                max1, _ = torch.max(max1, dim=1, keepdim=False)
                max2, _ = torch.max(embedding, dim=1, keepdim=False)
                max2, _ = torch.max(max2, dim=1, keepdim=False)
                sigmoided_scores = max1 + max2
                return sigmoided_scores

        model = TestModel()
        self.run_compare_torch([(1, 21, 2), (1, 6, 384)], model, backend=backend)# [spans.shape, embedding.shape]

class TestSlice(TorchBaseTest):
    @pytest.mark.skipif(_python_version() < (3, 6), reason="requires python 3.6")
    @pytest.mark.parametrize(
        "backend", backends,
    )
    def test_dynamic_slice(self, backend):
        class DynamicSlicer(torch.nn.Module):
            def __init__(self):
                super(DynamicSlicer, self).__init__()

            def forward(self, x, context_length):
                return x[context_length:, :, :]

        class Model(torch.nn.Module):

            def __init__(self):
                super(Model, self).__init__()
                self.tokens_embedding = torch.nn.Embedding(10, 10, 0)
                self.context_embedding = torch.nn.Embedding(10, 10, 0)
                self.dynamic_slicer = DynamicSlicer()

            def forward(self, tokens, context, context_length):
                tokens_embeddings = self.tokens_embedding(tokens)
                context_embeddings = self.context_embedding(context)
                embeddings = torch.cat((context_embeddings, tokens_embeddings), dim=0)
                embeddings = self.dynamic_slicer(embeddings, context_length)

                return embeddings

        model = Model()
        batch_size = 5
        inputs = [ TensorType(name="tokens", shape=(10, batch_size), dtype=np.int64),
                   TensorType(name="context", shape=(3, batch_size), dtype=np.int64),
                   TensorType(name="context_length", shape=(), dtype=np.int32),
                   ]
        self.run_compare_torch(inputs, model, rand_range=(0, 8),
                               backend=backend, use_scripting=False)

class TestRepeat(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, rank",
        itertools.product(backends, list(range(1, 6))),
    )
    def test_repeat(self, backend, rank):
        input_shape = np.random.randint(low=2, high=6, size=rank)
        repeats = np.random.randint(low=2, high=4, size=rank)
        input_shape = tuple(input_shape)

        model = ModuleWrapper(function=lambda x: x.repeat(*repeats))
        self.run_compare_torch(input_shape, model, backend=backend)

class TestStd(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, unbiased",
        itertools.product(backends, [True, False]),
    )
    def test_std_2_inputs(self, backend, unbiased):
        model = ModuleWrapper(function=torch.std,
                              kwargs={"unbiased": unbiased})
        x = torch.randn(1, 5, 10) * 3
        out = torch.std(x, unbiased=unbiased).unsqueeze(0)
        self.run_compare_torch(x, model, expected_results=out,
                           input_as_shape=False, backend=backend)

    @pytest.mark.parametrize(
        "backend, unbiased, dim, keepdim",
        itertools.product(backends, [True, False], [[0,2], [1], [2]], [True, False]),
    )
    def test_std_4_inputs(self, backend, unbiased, dim, keepdim):
        model = ModuleWrapper(function=torch.std,
                              kwargs={"unbiased": unbiased, "dim" : dim, "keepdim": keepdim})
        input_shape = (2, 5, 10)
        self.run_compare_torch(input_shape, model, backend=backend)

class TestTopk(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, largest, shape_dim_k",
        itertools.product(
            backends,
            [True, False],
            [
             ((4, 6, 7, 3), -1, 2),
             ((10, 3, 4), 2, 2),
             ((10, 5), -2, 3),
             ((5,), 0, 2)
             ],
        ),
    )
    def test_topk(self, backend, largest, shape_dim_k):
        input_shape = shape_dim_k[0]
        dim = shape_dim_k[1]
        k = shape_dim_k[2]

        class TopkModel(nn.Module):
            def __init__(self):
                super(TopkModel, self).__init__()

            def forward(self, x):
                return torch.topk(x, k, dim=dim, largest=largest)

        input_data = torch.rand(input_shape)
        model = TopkModel()
        expected_results = model(input_data)
        expected_results = [expected_results.values, expected_results.indices]
        self.run_compare_torch(
            input_data,
            model,
            expected_results=expected_results,
            input_as_shape=False,
            backend=backend,
        )

class TestLog10(TorchBaseTest):
    @pytest.mark.parametrize(
        "backend, rank", itertools.product(backends, range(1, 6)),
    )
    def test_log10(self, backend, rank):

        class Log10Model(nn.Module):
            def __init__(self):
                super(Log10Model, self).__init__()

            def forward(self, x):
                return torch.log10(x)

        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
        model = Log10Model()
        self.run_compare_torch(
            input_shape, model, backend=backend,
        )


text:
import copy
from typing import Dict, List, Optional, Tuple, Union

from matplotlib import pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
import numpy as np
from PIL import Image, ImageDraw

from .plot import show_img

def geojson_to_plt(geojson: List[List[int]],
                   imsize: Union[List, Tuple]
    ) -> Tuple[List[int], List[int]]:
    '''
    Transforms geojson polygon relative coordinates to matplotlib absolute
    coordinates ([x, y], [width, height])

    :param geojson: (list(list(float))) coordinates of polygon in geojson format
    :param imsize: tuple(int) height and width of an image
    :return: tuple(tuple(int, int)) rectangle coordinates in pixels in
        (top-left, bot-right) format
    '''
    geojson = copy.deepcopy(geojson)
    # scale by image size
    for point in geojson:
        point[0] = int(point[0] * imsize[1])  # horizontal coord
        point[1] = int((point[1]) * imsize[0])  # vertical coord

    # sort to find top left and bottom right corners
    geojson = sorted(geojson, key=lambda point: (point[0] + point[1]))
    top_left = geojson[0]
    bottom_right = geojson[-1]
    width = bottom_right[0] - top_left[0]
    height = bottom_right[1] - top_left[1]
    return (top_left, [width, height])

def plot_polygons(img: np.ndarray,
                  poly_dict: Dict,
                  alpha: float = 0.5,
                  ax: Optional[plt.Figure] = None,
                  figsize: Optional[Tuple[int, int]] = None,
                  show_axs: bool = False,
                  convert: bool = False):
    if not ax:
        _, ax = plt.subplots(figsize=figsize)
    show_img(img, ax, figsize, show_axs)
    patches = []
    colors = []

    labels = list(poly_dict.keys())
    for l in labels:
        color = 100*np.random.rand()
        for polygons in poly_dict[l]:
            for poly in polygons:
                if poly:
                    if convert:
                        poly = np.array(convert_poly(poly, img.shape))
                    patches.append(Polygon(poly, True))
            colors.append(color)
    patch_collection = PatchCollection(patches, alpha=alpha)
    patch_collection.set_array(np.array(colors))
    ax.add_collection(patch_collection)

def plot_bboxes(img: np.ndarray,
                bbox_dict: Dict,
                colormap: Dict,
                alpha: float = 0.5,
                ax: Optional[plt.Figure] = None,
                figsize: Optional[Tuple[int, int]] = None,
                show_axs: bool = False):
    if not ax:
        fig, ax = plt.subplots(figsize=figsize)
    show_img(img, ax, figsize, show_axs)
    patches = []
    colors = []

    labels = list(bbox_dict.keys())
    for l in labels:
        color = colormap[l]
        for poly in bbox_dict[l]:
            if poly:
                p = np.array(convert_poly(poly, img.shape))
                patches.append(Polygon(p, True))
            colors.append(color)
    patch_collection = PatchCollection(patches, alpha=alpha)
    patch_collection.set_array(np.array(colors))
    ax.add_collection(patch_collection)

def convert_poly(poly: List[float],
                 imshape: Union[List[int], Tuple[int, int]]) -> List[int]:
    return list(map(lambda p: (int(p[0]*imshape[1]), int(p[1]*imshape[0])),
                    poly))

def mask_from_polygon(polygon: List[int],
                      imshape: Union[List[int], Tuple[int, int]]
    ) -> np.ndarray:
    img = Image.new('L', (imshape[1], imshape[0]), 0)
    ImageDraw.Draw(img).polygon(polygon, outline=1, fill=1)
    return np.array(img)

ext: py

max_stars_repo_name: belskikh/plottify

max_issues_repo_name: belskikh/plottify

max_forks_repo_name: belskikh/plottify

reason: import numpy


text:
# Based on https://www.tensorflow.org/tutorials/images/transfer_learning

# Imports
import os
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
keras = tf.keras

import tensorflow_datasets as tfds
tfds.disable_progress_bar()

# Split the data manually into 80% training, 10% testing, 10% validation
(raw_train, raw_validation, raw_test), metadata = tfds.load(
    'cats_vs_dogs',
    split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'],
    with_info=True,
    as_supervised=True,
)

get_label_name = metadata.features['label'].int2str  # creates a function object that we can use to get labels

# display 5 images from the dataset
for image, label in raw_train.take(5):
	plt.figure()
  	plt.imshow(image)
  	plt.title(get_label_name(label))

IMG_SIZE = 160 # All images will be resized to 160x160

def format_example(image, label):
	#returns an image that is reshaped to IMG_SIZE
	image = tf.cast(image, tf.float32)
	image = (image/127.5) - 1
	image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
	return image, label

# Applying the function to the whole dataset
train = raw_train.map(format_example)
validation = raw_validation.map(format_example)
test = raw_test.map(format_example)

# Checking the images. (Images do look darker, investigate)
for image, label in train.take(2):
	plt.figure()
	plt.imshow(image)
	plt.title(get_label_name(label))

# Shuffling and splitting the images in batches
BATCH_SIZE = 32
SHUFFLE_BUFFER_SIZE = 1000

train_batches = train.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE)
validation_batches = validation.batch(BATCH_SIZE)
test_batches = test.batch(BATCH_SIZE)

# Checking the image size
for img, label in raw_train.take(2):
	print("Original shape:", img.shape)

for img, label in train.take(2):
	print("New shape:", img.shape)

# Picking a pretrained model:
IMG_SHAPE = (IMG_SIZE, IMG_SIZE, 3)

# Create the base model from the pre-trained model MobileNet V2
base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
                                               include_top=False,  # We do not want the top (classification) layer!
                                               weights='imagenet')

# This is what the pretrained model looks like:
base_model.summary()
# Model: "mobilenetv2_1.00_160"
# __________________________________________________________________________________________________
# Layer (type)                    Output Shape         Param #     Connected to
# ==================================================================================================
# input_1 (InputLayer)            [(None, 160, 160, 3) 0
# __________________________________________________________________________________________________
# Conv1 (Conv2D)                  (None, 80, 80, 32)   864         input_1[0][0]
# __________________________________________________________________________________________________
# bn_Conv1 (BatchNormalization)   (None, 80, 80, 32)   128         Conv1[0][0]
# __________________________________________________________________________________________________
# Conv1_relu (ReLU)               (None, 80, 80, 32)   0           bn_Conv1[0][0]
# __________________________________________________________________________________________________
# expanded_conv_depthwise (Depthw (None, 80, 80, 32)   288         Conv1_relu[0][0]
# __________________________________________________________________________________________________
# expanded_conv_depthwise_BN (Bat (None, 80, 80, 32)   128         expanded_conv_depthwise[0][0]
# __________________________________________________________________________________________________
# expanded_conv_depthwise_relu (R (None, 80, 80, 32)   0           expanded_conv_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# expanded_conv_project (Conv2D)  (None, 80, 80, 16)   512         expanded_conv_depthwise_relu[0][0
# __________________________________________________________________________________________________
# expanded_conv_project_BN (Batch (None, 80, 80, 16)   64          expanded_conv_project[0][0]
# __________________________________________________________________________________________________
# block_1_expand (Conv2D)         (None, 80, 80, 96)   1536        expanded_conv_project_BN[0][0]
# __________________________________________________________________________________________________
# block_1_expand_BN (BatchNormali (None, 80, 80, 96)   384         block_1_expand[0][0]
# __________________________________________________________________________________________________
# block_1_expand_relu (ReLU)      (None, 80, 80, 96)   0           block_1_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_1_pad (ZeroPadding2D)     (None, 81, 81, 96)   0           block_1_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_1_depthwise (DepthwiseCon (None, 40, 40, 96)   864         block_1_pad[0][0]
# __________________________________________________________________________________________________
# block_1_depthwise_BN (BatchNorm (None, 40, 40, 96)   384         block_1_depthwise[0][0]
# __________________________________________________________________________________________________
# block_1_depthwise_relu (ReLU)   (None, 40, 40, 96)   0           block_1_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_1_project (Conv2D)        (None, 40, 40, 24)   2304        block_1_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_1_project_BN (BatchNormal (None, 40, 40, 24)   96          block_1_project[0][0]
# __________________________________________________________________________________________________
# block_2_expand (Conv2D)         (None, 40, 40, 144)  3456        block_1_project_BN[0][0]
# __________________________________________________________________________________________________
# block_2_expand_BN (BatchNormali (None, 40, 40, 144)  576         block_2_expand[0][0]
# __________________________________________________________________________________________________
# block_2_expand_relu (ReLU)      (None, 40, 40, 144)  0           block_2_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_2_depthwise (DepthwiseCon (None, 40, 40, 144)  1296        block_2_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_2_depthwise_BN (BatchNorm (None, 40, 40, 144)  576         block_2_depthwise[0][0]
# __________________________________________________________________________________________________
# block_2_depthwise_relu (ReLU)   (None, 40, 40, 144)  0           block_2_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_2_project (Conv2D)        (None, 40, 40, 24)   3456        block_2_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_2_project_BN (BatchNormal (None, 40, 40, 24)   96          block_2_project[0][0]
# __________________________________________________________________________________________________
# block_2_add (Add)               (None, 40, 40, 24)   0           block_1_project_BN[0][0]
#                                                                  block_2_project_BN[0][0]
# __________________________________________________________________________________________________
# block_3_expand (Conv2D)         (None, 40, 40, 144)  3456        block_2_add[0][0]
# __________________________________________________________________________________________________
# block_3_expand_BN (BatchNormali (None, 40, 40, 144)  576         block_3_expand[0][0]
# __________________________________________________________________________________________________
# block_3_expand_relu (ReLU)      (None, 40, 40, 144)  0           block_3_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_3_pad (ZeroPadding2D)     (None, 41, 41, 144)  0           block_3_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_3_depthwise (DepthwiseCon (None, 20, 20, 144)  1296        block_3_pad[0][0]
# __________________________________________________________________________________________________
# block_3_depthwise_BN (BatchNorm (None, 20, 20, 144)  576         block_3_depthwise[0][0]
# __________________________________________________________________________________________________
# block_3_depthwise_relu (ReLU)   (None, 20, 20, 144)  0           block_3_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_3_project (Conv2D)        (None, 20, 20, 32)   4608        block_3_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_3_project_BN (BatchNormal (None, 20, 20, 32)   128         block_3_project[0][0]
# __________________________________________________________________________________________________
# block_4_expand (Conv2D)         (None, 20, 20, 192)  6144        block_3_project_BN[0][0]
# __________________________________________________________________________________________________
# block_4_expand_BN (BatchNormali (None, 20, 20, 192)  768         block_4_expand[0][0]
# __________________________________________________________________________________________________
# block_4_expand_relu (ReLU)      (None, 20, 20, 192)  0           block_4_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_4_depthwise (DepthwiseCon (None, 20, 20, 192)  1728        block_4_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_4_depthwise_BN (BatchNorm (None, 20, 20, 192)  768         block_4_depthwise[0][0]
# __________________________________________________________________________________________________
# block_4_depthwise_relu (ReLU)   (None, 20, 20, 192)  0           block_4_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_4_project (Conv2D)        (None, 20, 20, 32)   6144        block_4_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_4_project_BN (BatchNormal (None, 20, 20, 32)   128         block_4_project[0][0]
# __________________________________________________________________________________________________
# block_4_add (Add)               (None, 20, 20, 32)   0           block_3_project_BN[0][0]
#                                                                  block_4_project_BN[0][0]
# __________________________________________________________________________________________________
# block_5_expand (Conv2D)         (None, 20, 20, 192)  6144        block_4_add[0][0]
# __________________________________________________________________________________________________
# block_5_expand_BN (BatchNormali (None, 20, 20, 192)  768         block_5_expand[0][0]
# __________________________________________________________________________________________________
# block_5_expand_relu (ReLU)      (None, 20, 20, 192)  0           block_5_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_5_depthwise (DepthwiseCon (None, 20, 20, 192)  1728        block_5_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_5_depthwise_BN (BatchNorm (None, 20, 20, 192)  768         block_5_depthwise[0][0]
# __________________________________________________________________________________________________
# block_5_depthwise_relu (ReLU)   (None, 20, 20, 192)  0           block_5_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_5_project (Conv2D)        (None, 20, 20, 32)   6144        block_5_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_5_project_BN (BatchNormal (None, 20, 20, 32)   128         block_5_project[0][0]
# __________________________________________________________________________________________________
# block_5_add (Add)               (None, 20, 20, 32)   0           block_4_add[0][0]
#                                                                  block_5_project_BN[0][0]
# __________________________________________________________________________________________________
# block_6_expand (Conv2D)         (None, 20, 20, 192)  6144        block_5_add[0][0]
# __________________________________________________________________________________________________
# block_6_expand_BN (BatchNormali (None, 20, 20, 192)  768         block_6_expand[0][0]
# __________________________________________________________________________________________________
# block_6_expand_relu (ReLU)      (None, 20, 20, 192)  0           block_6_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_6_pad (ZeroPadding2D)     (None, 21, 21, 192)  0           block_6_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_6_depthwise (DepthwiseCon (None, 10, 10, 192)  1728        block_6_pad[0][0]
# __________________________________________________________________________________________________
# block_6_depthwise_BN (BatchNorm (None, 10, 10, 192)  768         block_6_depthwise[0][0]
# __________________________________________________________________________________________________
# block_6_depthwise_relu (ReLU)   (None, 10, 10, 192)  0           block_6_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_6_project (Conv2D)        (None, 10, 10, 64)   12288       block_6_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_6_project_BN (BatchNormal (None, 10, 10, 64)   256         block_6_project[0][0]
# __________________________________________________________________________________________________
# block_7_expand (Conv2D)         (None, 10, 10, 384)  24576       block_6_project_BN[0][0]
# __________________________________________________________________________________________________
# block_7_expand_BN (BatchNormali (None, 10, 10, 384)  1536        block_7_expand[0][0]
# __________________________________________________________________________________________________
# block_7_expand_relu (ReLU)      (None, 10, 10, 384)  0           block_7_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_7_depthwise (DepthwiseCon (None, 10, 10, 384)  3456        block_7_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_7_depthwise_BN (BatchNorm (None, 10, 10, 384)  1536        block_7_depthwise[0][0]
# __________________________________________________________________________________________________
# block_7_depthwise_relu (ReLU)   (None, 10, 10, 384)  0           block_7_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_7_project (Conv2D)        (None, 10, 10, 64)   24576       block_7_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_7_project_BN (BatchNormal (None, 10, 10, 64)   256         block_7_project[0][0]
# __________________________________________________________________________________________________
# block_7_add (Add)               (None, 10, 10, 64)   0           block_6_project_BN[0][0]
#                                                                  block_7_project_BN[0][0]
# __________________________________________________________________________________________________
# block_8_expand (Conv2D)         (None, 10, 10, 384)  24576       block_7_add[0][0]
# __________________________________________________________________________________________________
# block_8_expand_BN (BatchNormali (None, 10, 10, 384)  1536        block_8_expand[0][0]
# __________________________________________________________________________________________________
# block_8_expand_relu (ReLU)      (None, 10, 10, 384)  0           block_8_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_8_depthwise (DepthwiseCon (None, 10, 10, 384)  3456        block_8_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_8_depthwise_BN (BatchNorm (None, 10, 10, 384)  1536        block_8_depthwise[0][0]
# __________________________________________________________________________________________________
# block_8_depthwise_relu (ReLU)   (None, 10, 10, 384)  0           block_8_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_8_project (Conv2D)        (None, 10, 10, 64)   24576       block_8_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_8_project_BN (BatchNormal (None, 10, 10, 64)   256         block_8_project[0][0]
# __________________________________________________________________________________________________
# block_8_add (Add)               (None, 10, 10, 64)   0           block_7_add[0][0]
#                                                                  block_8_project_BN[0][0]
# __________________________________________________________________________________________________
# block_9_expand (Conv2D)         (None, 10, 10, 384)  24576       block_8_add[0][0]
# __________________________________________________________________________________________________
# block_9_expand_BN (BatchNormali (None, 10, 10, 384)  1536        block_9_expand[0][0]
# __________________________________________________________________________________________________
# block_9_expand_relu (ReLU)      (None, 10, 10, 384)  0           block_9_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_9_depthwise (DepthwiseCon (None, 10, 10, 384)  3456        block_9_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_9_depthwise_BN (BatchNorm (None, 10, 10, 384)  1536        block_9_depthwise[0][0]
# __________________________________________________________________________________________________
# block_9_depthwise_relu (ReLU)   (None, 10, 10, 384)  0           block_9_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_9_project (Conv2D)        (None, 10, 10, 64)   24576       block_9_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_9_project_BN (BatchNormal (None, 10, 10, 64)   256         block_9_project[0][0]
# __________________________________________________________________________________________________
# block_9_add (Add)               (None, 10, 10, 64)   0           block_8_add[0][0]
#                                                                  block_9_project_BN[0][0]
# __________________________________________________________________________________________________
# block_10_expand (Conv2D)        (None, 10, 10, 384)  24576       block_9_add[0][0]
# __________________________________________________________________________________________________
# block_10_expand_BN (BatchNormal (None, 10, 10, 384)  1536        block_10_expand[0][0]
# __________________________________________________________________________________________________
# block_10_expand_relu (ReLU)     (None, 10, 10, 384)  0           block_10_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_10_depthwise (DepthwiseCo (None, 10, 10, 384)  3456        block_10_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_10_depthwise_BN (BatchNor (None, 10, 10, 384)  1536        block_10_depthwise[0][0]
# __________________________________________________________________________________________________
# block_10_depthwise_relu (ReLU)  (None, 10, 10, 384)  0           block_10_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_10_project (Conv2D)       (None, 10, 10, 96)   36864       block_10_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_10_project_BN (BatchNorma (None, 10, 10, 96)   384         block_10_project[0][0]
# __________________________________________________________________________________________________
# block_11_expand (Conv2D)        (None, 10, 10, 576)  55296       block_10_project_BN[0][0]
# __________________________________________________________________________________________________
# block_11_expand_BN (BatchNormal (None, 10, 10, 576)  2304        block_11_expand[0][0]
# __________________________________________________________________________________________________
# block_11_expand_relu (ReLU)     (None, 10, 10, 576)  0           block_11_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_11_depthwise (DepthwiseCo (None, 10, 10, 576)  5184        block_11_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_11_depthwise_BN (BatchNor (None, 10, 10, 576)  2304        block_11_depthwise[0][0]
# __________________________________________________________________________________________________
# block_11_depthwise_relu (ReLU)  (None, 10, 10, 576)  0           block_11_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_11_project (Conv2D)       (None, 10, 10, 96)   55296       block_11_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_11_project_BN (BatchNorma (None, 10, 10, 96)   384         block_11_project[0][0]
# __________________________________________________________________________________________________
# block_11_add (Add)              (None, 10, 10, 96)   0           block_10_project_BN[0][0]
#                                                                  block_11_project_BN[0][0]
# __________________________________________________________________________________________________
# block_12_expand (Conv2D)        (None, 10, 10, 576)  55296       block_11_add[0][0]
# __________________________________________________________________________________________________
# block_12_expand_BN (BatchNormal (None, 10, 10, 576)  2304        block_12_expand[0][0]
# __________________________________________________________________________________________________
# block_12_expand_relu (ReLU)     (None, 10, 10, 576)  0           block_12_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_12_depthwise (DepthwiseCo (None, 10, 10, 576)  5184        block_12_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_12_depthwise_BN (BatchNor (None, 10, 10, 576)  2304        block_12_depthwise[0][0]
# __________________________________________________________________________________________________
# block_12_depthwise_relu (ReLU)  (None, 10, 10, 576)  0           block_12_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_12_project (Conv2D)       (None, 10, 10, 96)   55296       block_12_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_12_project_BN (BatchNorma (None, 10, 10, 96)   384         block_12_project[0][0]
# __________________________________________________________________________________________________
# block_12_add (Add)              (None, 10, 10, 96)   0           block_11_add[0][0]
#                                                                  block_12_project_BN[0][0]
# __________________________________________________________________________________________________
# block_13_expand (Conv2D)        (None, 10, 10, 576)  55296       block_12_add[0][0]
# __________________________________________________________________________________________________
# block_13_expand_BN (BatchNormal (None, 10, 10, 576)  2304        block_13_expand[0][0]
# __________________________________________________________________________________________________
# block_13_expand_relu (ReLU)     (None, 10, 10, 576)  0           block_13_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_13_pad (ZeroPadding2D)    (None, 11, 11, 576)  0           block_13_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_13_depthwise (DepthwiseCo (None, 5, 5, 576)    5184        block_13_pad[0][0]
# __________________________________________________________________________________________________
# block_13_depthwise_BN (BatchNor (None, 5, 5, 576)    2304        block_13_depthwise[0][0]
# __________________________________________________________________________________________________
# block_13_depthwise_relu (ReLU)  (None, 5, 5, 576)    0           block_13_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_13_project (Conv2D)       (None, 5, 5, 160)    92160       block_13_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_13_project_BN (BatchNorma (None, 5, 5, 160)    640         block_13_project[0][0]
# __________________________________________________________________________________________________
# block_14_expand (Conv2D)        (None, 5, 5, 960)    153600      block_13_project_BN[0][0]
# __________________________________________________________________________________________________
# block_14_expand_BN (BatchNormal (None, 5, 5, 960)    3840        block_14_expand[0][0]
# __________________________________________________________________________________________________
# block_14_expand_relu (ReLU)     (None, 5, 5, 960)    0           block_14_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_14_depthwise (DepthwiseCo (None, 5, 5, 960)    8640        block_14_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_14_depthwise_BN (BatchNor (None, 5, 5, 960)    3840        block_14_depthwise[0][0]
# __________________________________________________________________________________________________
# block_14_depthwise_relu (ReLU)  (None, 5, 5, 960)    0           block_14_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_14_project (Conv2D)       (None, 5, 5, 160)    153600      block_14_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_14_project_BN (BatchNorma (None, 5, 5, 160)    640         block_14_project[0][0]
# __________________________________________________________________________________________________
# block_14_add (Add)              (None, 5, 5, 160)    0           block_13_project_BN[0][0]
#                                                                  block_14_project_BN[0][0]
# __________________________________________________________________________________________________
# block_15_expand (Conv2D)        (None, 5, 5, 960)    153600      block_14_add[0][0]
# __________________________________________________________________________________________________
# block_15_expand_BN (BatchNormal (None, 5, 5, 960)    3840        block_15_expand[0][0]
# __________________________________________________________________________________________________
# block_15_expand_relu (ReLU)     (None, 5, 5, 960)    0           block_15_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_15_depthwise (DepthwiseCo (None, 5, 5, 960)    8640        block_15_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_15_depthwise_BN (BatchNor (None, 5, 5, 960)    3840        block_15_depthwise[0][0]
# __________________________________________________________________________________________________
# block_15_depthwise_relu (ReLU)  (None, 5, 5, 960)    0           block_15_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_15_project (Conv2D)       (None, 5, 5, 160)    153600      block_15_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_15_project_BN (BatchNorma (None, 5, 5, 160)    640         block_15_project[0][0]
# __________________________________________________________________________________________________
# block_15_add (Add)              (None, 5, 5, 160)    0           block_14_add[0][0]
#                                                                  block_15_project_BN[0][0]
# __________________________________________________________________________________________________
# block_16_expand (Conv2D)        (None, 5, 5, 960)    153600      block_15_add[0][0]
# __________________________________________________________________________________________________
# block_16_expand_BN (BatchNormal (None, 5, 5, 960)    3840        block_16_expand[0][0]
# __________________________________________________________________________________________________
# block_16_expand_relu (ReLU)     (None, 5, 5, 960)    0           block_16_expand_BN[0][0]
# __________________________________________________________________________________________________
# block_16_depthwise (DepthwiseCo (None, 5, 5, 960)    8640        block_16_expand_relu[0][0]
# __________________________________________________________________________________________________
# block_16_depthwise_BN (BatchNor (None, 5, 5, 960)    3840        block_16_depthwise[0][0]
# __________________________________________________________________________________________________
# block_16_depthwise_relu (ReLU)  (None, 5, 5, 960)    0           block_16_depthwise_BN[0][0]
# __________________________________________________________________________________________________
# block_16_project (Conv2D)       (None, 5, 5, 320)    307200      block_16_depthwise_relu[0][0]
# __________________________________________________________________________________________________
# block_16_project_BN (BatchNorma (None, 5, 5, 320)    1280        block_16_project[0][0]
# __________________________________________________________________________________________________
# Conv_1 (Conv2D)                 (None, 5, 5, 1280)   409600      block_16_project_BN[0][0]
# __________________________________________________________________________________________________
# Conv_1_bn (BatchNormalization)  (None, 5, 5, 1280)   5120        Conv_1[0][0]
# __________________________________________________________________________________________________
# out_relu (ReLU)                 (None, 5, 5, 1280)   0           Conv_1_bn[0][0]
# ==================================================================================================
# Total params: 2,257,984
# Trainable params: 2,223,872
# Non-trainable params: 34,112
# __________________________________________________________________________________________________

# We can see there are over 2.2M trainable params at the moment! We need to reduce that soon!

# At this point this base_model will simply output a shape (32, 5, 5, 1280) tensor that is a feature extraction
# from our original (1, 160, 160, 3) image. The 32 means that we have 32 layers of different filters/features.
for image, _ in train_batches.take(1):
   pass

feature_batch = base_model(image)
print(feature_batch.shape)

# We will make the model untrainable, as it is already trained
base_model.trainable = False
# If we use base_model.summary() again, we will see the trainable params is 0

# We now create our layers:
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
prediction_layer = keras.layers.Dense(1)

# And combine everything together into a model with Keras:
model = tf.keras.Sequential([
  base_model,
  global_average_layer,
  prediction_layer
])

# This is our full model now:
model.summary()
# Model: "sequential_2"
# _________________________________________________________________
# Layer (type)                 Output Shape              Param #
# =================================================================
# mobilenetv2_1.00_160 (Functi (None, 5, 5, 1280)        2257984
# _________________________________________________________________
# global_average_pooling2d (Gl (None, 1280)              0
# _________________________________________________________________
# dense_4 (Dense)              (None, 1)                 1281
# =================================================================
# Total params: 2,259,265
# Trainable params: 1,281
# Non-trainable params: 2,257,984
# _________________________________________________________________

# The trainable params are only now the layers that we added on top of the base model, much less than before

# Training the model:
base_learning_rate = 0.0001  # Very small learning rate to ensure that there are not really major changes to the model
model.compile(optimizer=tf.keras.optimizers.RMSprop(lr=base_learning_rate),
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['accuracy'])

# We can evaluate the model right now to see how it does before training it on our new images
initial_epochs = 3
validation_steps=20

loss0,accuracy0 = model.evaluate(validation_batches, steps = validation_steps)

# Now we can train it on our images
history = model.fit(train_batches,
                    epochs=initial_epochs,
                    validation_data=validation_batches)

# This accuracy is much higher now, almost 98%!
acc = history.history['accuracy']
print(acc)

# We can also save and load models using Keras:
model.save("dogs_vs_cats.h5")  # we can save the model and reload it at anytime in the future
new_model = tf.keras.models.load_model("dogs_vs_cats.h5")

max_stars_repo_path: 5 - Computer Vision/image_detection_with_pretrained_model.py


text:
import unittest
import numpy as np
from numpy import pi

from robot_arm import RobotArm

class TestRobotArm(unittest.TestCase):
    def setUp(self):
        self.lengths = (3, 2, 2,)
        self.destinations = (
            (5, 4, 6, 4, 5),
            (0, 2, 0.5, -2, -1),
        )
        self.theta = (pi, pi / 2, 0,)
        self.robot_arm = RobotArm(self.lengths, self.destinations, self.theta)

    def test_init_all_arguments(self):
        RobotArm(self.lengths, self.destinations, self.theta)

    def test_init_without_theta(self):
        RobotArm(self.lengths, self.destinations)

    def test_wrong_lengths_type(self):
        self.assertRaises(
            TypeError,
            RobotArm,
            np.array(self.lengths),
            self.destinations,
            self.theta
        )

    def test_wrong_destinations_type(self):
        self.assertRaises(
            TypeError,
            RobotArm,
            self.lengths,
            np.array(self.destinations),
            self.theta)

    def test_wrong_theta_type(self):
        self.assertRaises(
            TypeError,
            RobotArm,
            self.lengths,
            self.destinations,
            np.array(self.theta))

    def test_destinations_properties(self):
        robot_arm = RobotArm(self.lengths, self.destinations, self.theta)
        self.assertIsInstance(robot_arm.destinations, np.ndarray)

        # Check if points are 2D
        self.assertTrue(robot_arm.destinations.shape[0] == 2)

        # Check if destinations are immutable
        self.assertRaises(
            ValueError,
            robot_arm.destinations.__setitem__,
            (0, 0,),
            0,
        )

    def test_generate_initial_guess(self):
        self.robot_arm.generate_initial_guess()

max_stars_repo_name: JakobGM/robotarm-optimization

max_issues_repo_name: JakobGM/robotarm-optimization

max_forks_repo_name: JakobGM/robotarm-optimization

reason: import numpy,from numpy


text:
from portfolio import Portfolio
import numpy as np

# Just a wrapper to use the Portfolio class without caring about allocation and direction vectors
class Universe(Portfolio):

    def __init__(self, assets, riskFree):
        allocation = np.ones(len(assets))
        direction = np.ones(len(assets))
        super().__init__(assets, allocation, direction, riskFree)

max_stars_repo_name: iamthejao/portfolioTools

max_issues_repo_name: iamthejao/portfolioTools

max_forks_repo_name: iamthejao/portfolioTools


text:
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn import Sequential
import matplotlib.pyplot as plt

n_hidden = 4
nn1 = Sequential(nn.Linear(2, n_hidden), nn.Tanh(), nn.Linear(n_hidden, 2))
nn2 = Sequential(nn.Linear(2, n_hidden), nn.Tanh(),
                  nn.Linear(n_hidden, n_hidden), nn.Tanh(),
                  nn.Linear(n_hidden, 2))

n = 1000
n_points = 10
t = np.linspace(0, 2 * np.pi, n_points, endpoint=False)
c = np.array([(np.cos(t_), np.sin(t_)) for t_ in t])
y = np.arange(n_points) % 2

X = np.concatenate([0.1 * np.random.randn(n, 2) + c_ for c_ in c])
y = np.concatenate([y_ * np.ones(n) for y_ in y])
X = torch.tensor(X).float()
y = torch.tensor(y).long()

f, ax = plt.subplots(figsize=(3.5, 3.5))
xm, xM = -1.5, 1.5
ax.set_xlim(xm, xM)
ax.set_ylim(xm, xM)
s = 3
for i, name in enumerate(['class 1', 'class 2']):
    loc = np.where(y == i)[0]
    plt.scatter(X[loc, 0], X[loc, 1], s=s, label=name)
plt.legend()
ax.set_xticks([])
ax.set_yticks([])

plt.savefig('images/nn.png', dpi=200)

for n_hidden in [2, 3, 4, 5, 6]:
    nn1 = Sequential(nn.Linear(2, n_hidden), nn.Tanh(),
                     nn.Linear(n_hidden, 2 * n_hidden), nn.Tanh(),
                     nn.Linear(2 * n_hidden, 2))
    optimizer = optim.Adam(nn1.parameters(), lr=1e-2)
    criterion = nn.CrossEntropyLoss()
    for i in range(1001):
        optimizer.zero_grad()
        pred = nn1(X)
        loss = criterion(pred, y)
        loss.backward()
        optimizer.step()
        if i % 100 == 0:
            print(loss.item())

    f, ax = plt.subplots(figsize=(3.5, 3.5))
    xm, xM = -1.5, 1.5
    ax.set_xlim(xm, xM)
    ax.set_ylim(xm, xM)
    s = 3
    for i, name in enumerate(['class 1', 'class 2']):
        loc = np.where(y == i)[0]
        plt.scatter(X[loc, 0], X[loc, 1], s=s, label=name)
    plt.legend()
    ax.set_xticks([])
    ax.set_yticks([])

    xx, yy = np.meshgrid(np.linspace(-1.5, 1.5),
                         np.linspace(-1.5, 1.5))
    data = torch.tensor(np.c_[xx.ravel(), yy.ravel()]).float()
    op = nn1(data).detach()
    z = op.numpy().argmax(axis=1)
    Z = z.reshape(xx.shape)
    plt.contourf(xx, yy, Z, levels=1, alpha=0.5, colors=['b', 'orange'])
    plt.savefig('images/nn_two_%s.png' % n_hidden, dpi=200)

for n_hidden in [2, 3, 4, 5, 6]:
    nn1 = Sequential(nn.Linear(2, n_hidden), nn.Tanh(), nn.Linear(n_hidden, 2))
    optimizer = optim.Adam(nn1.parameters(), lr=1e-2)
    criterion = nn.CrossEntropyLoss()
    for i in range(1001):
        optimizer.zero_grad()
        pred = nn1(X)
        loss = criterion(pred, y)
        loss.backward()
        optimizer.step()
        if i % 100 == 0:
            print(loss.item())

    f, ax = plt.subplots(figsize=(3.5, 3.5))
    xm, xM = -1.5, 1.5
    ax.set_xlim(xm, xM)
    ax.set_ylim(xm, xM)
    s = 3
    for i, name in enumerate(['class 1', 'class 2']):
        loc = np.where(y == i)[0]
        plt.scatter(X[loc, 0], X[loc, 1], s=s, label=name)
    plt.legend()
    ax.set_xticks([])
    ax.set_yticks([])

    xx, yy = np.meshgrid(np.linspace(-1.5, 1.5),
                         np.linspace(-1.5, 1.5))
    data = torch.tensor(np.c_[xx.ravel(), yy.ravel()]).float()
    op = nn1(data).detach()
    z = op.numpy().argmax(axis=1)
    Z = z.reshape(xx.shape)
    plt.contourf(xx, yy, Z, levels=1, alpha=0.5, colors=['b', 'orange'])
    plt.savefig('images/nn_one_%s.png' % n_hidden, dpi=200)