"""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, 0, 8, 20), 'torch.manual_seed', 'torch.manual_seed', ({(8, 18, 8, 19): '(0)'}, {}), '(0)', False, 'import torch\n'), ((11, 9, 11, 37), 'torch.randn', 'torch.randn', ({(11, 21, 11, 36): '(64, 3, 32, 32)'}, {}), '((64, 3, 32, 32))', False, 'import torch\n'), ((20, 15, 22, 1), 'torch.nn.Unfold', 'torch.nn.Unfold', (), '', False, 'import torch\n'), ((23, 13, 25, 1), 'unfoldNd.UnfoldNd', 'unfoldNd.UnfoldNd', (), '', False, 'import unfoldNd\n'), ((32, 3, 32, 45), 'torch.allclose', 'torch.allclose', ({(32, 18, 32, 31): 'torch_outputs', (32, 33, 32, 44): 'lib_outputs'}, {}), '(torch_outputs, lib_outputs)', False, 'import torch\n')]