feat: simplify, add skip items, resnet groups

README.md

@@ -1,6 +1,6 @@
 # A-UNet
 
-A library that provides building blocks to customize UNets, in PyTorch.
+A toolbox that provides hackable building blocks for (1D/2D/3D) UNets, in PyTorch.
 
 ## Install
 ```bash

a_unet/blocks.py

@@ -25,7 +25,7 @@ def __call__(self, *a, **ka):
         class Inner:
             def __init__(self):
                 self.args = a
-                self.kwargs = ka
+                self.__dict__.update(**ka)
 
             def __call__(self, *b, **kb):
                 if override:
@@ -320,11 +320,11 @@ def forward(x: Tensor, features: Tensor) -> Tensor:
     return Module([to_scale_shift, norm], forward)
 
 
-def MergeAdd(*args, **kwargs):
+def MergeAdd():
     return Module([], lambda x, y, *_: x + y)
 
 
-def MergeCat(dim: int, channels: int, **kwargs) -> nn.Module:
+def MergeCat(dim: int, channels: int) -> nn.Module:
     conv = Conv(dim=dim, in_channels=channels * 2, out_channels=channels, kernel_size=1)
     return Module([conv], lambda x, y, *_: conv(torch.cat([x, y], dim=1)))