"""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")

[((145, 165), 'torch.manual_seed', 'torch.manual_seed', (['(0)'], {}), '(0)\n', (162, 165), False, 'import torch\n'), ((246, 274), 'torch.randn', 'torch.randn', (['(64, 3, 32, 32)'], {}), '((64, 3, 32, 32))\n', (257, 274), False, 'import torch\n'), ((441, 520), 'torch.nn.Unfold', 'torch.nn.Unfold', (['kernel_size'], {'dilation': 'dilation', 'padding': 'padding', 'stride': 'stride'}), '(kernel_size, dilation=dilation, padding=padding, stride=stride)\n', (456, 520), False, 'import torch\n'), ((540, 626), 'unfoldNd.UnfoldNd', 'unfoldNd.UnfoldNd', (['kernel_size'], {'dilation': 'dilation', 'padding': 'padding', 'stride': 'stride'}), '(kernel_size, dilation=dilation, padding=padding, stride=\n    stride)\n', (557, 626), False, 'import unfoldNd\n'), ((726, 768), 'torch.allclose', 'torch.allclose', (['torch_outputs', 'lib_outputs'], {}), '(torch_outputs, lib_outputs)\n', (740, 768), False, 'import torch\n')]