"""How to use ``unfoldNd``. A comparison with ``torch.nn.Unfold``."""

# imports, make this example deterministic
import torch

import unfoldNd

torch.manual_seed(0)

# random batched RGB 32x32 image-shaped input tensor of batch size 64
inputs = torch.randn((64, 3, 32, 32))

# module hyperparameters
kernel_size = 3
dilation = 1
padding = 1
stride = 2

# both modules accept the same arguments and perform the same operation
torch_module = torch.nn.Unfold(
    kernel_size, dilation=dilation, padding=padding, stride=stride
)
lib_module = unfoldNd.UnfoldNd(
    kernel_size, dilation=dilation, padding=padding, stride=stride
)

# forward pass
torch_outputs = torch_module(inputs)
lib_outputs = lib_module(inputs)

# check
if torch.allclose(torch_outputs, lib_outputs):
    print("✔ Outputs of torch.nn.Unfold and unfoldNd.UnfoldNd match.")
else:
    raise AssertionError("❌ Outputs don't match")

[(8, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (11, 'torch.randn', 'torch.randn', False, 'import torch\n'), (20, 'torch.nn.Unfold', 'torch.nn.Unfold', False, 'import torch\n'), (23, 'unfoldNd.UnfoldNd', 'unfoldNd.UnfoldNd', False, 'import unfoldNd\n'), (32, 'torch.allclose', 'torch.allclose', False, 'import torch\n')]