Fix bug in timm.layers.drop.drop_block_2d when H != W.

There are two bugs in the `valid_block` code for `drop_block_2d`.
- a (W, H) grid being reshaped as (H, W)

The current code uses (W, H) to generate the meshgrid;
but then uses a `.reshape((1, 1, H, W))` to unsqueeze the block map.

The simplest fix to the first bug is a one-line change:
```python
h_i, w_i = ndgrid(torch.arange(H), torch.arange(W))
```

This is a longer patch, that attempts to make the code testable.

Note: The current code behaves oddly when the block_size or
clipped_block_size is even; I've added tests exposing the behavior;
but have not changed it.

When you trigger the reshape bug, you get wild results:
```
$ python scratch.py
{'H': 4, 'W': 5, 'block_size': 3, 'fix_reshape': False}
grid.shape=torch.Size([1, 1, 4, 5])
tensor([[[[False, False, False, False, False],
          [ True,  True, False, False,  True],
          [ True, False, False,  True,  True],
          [False, False, False, False, False]]]])

{'H': 4, 'W': 5, 'block_size': 3, 'fix_reshape': True}
grid.shape=torch.Size([1, 1, 4, 5])
tensor([[[[False, False, False, False, False],
          [False,  True,  True,  True, False],
          [False,  True,  True,  True, False],
          [False, False, False, False, False]]]])
```

Here's a tiny exceprt script, showing the problem;
it generated the above output.

```python
import torch
from typing import Tuple

def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]:
    """generate N-D grid in dimension order.

    The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.

    That is, the statement
    [X1,X2,X3] = ndgrid(x1,x2,x3)

    produces the same result as

    [X2,X1,X3] = meshgrid(x2,x1,x3)

    This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make
    torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy').

    """
    try:
        return torch.meshgrid(*tensors, indexing='ij')
    except TypeError:
        # old PyTorch < 1.10 will follow this path as it does not have indexing arg,
        # the old behaviour of meshgrid was 'ij'
        return torch.meshgrid(*tensors)

def valid_block(H, W, block_size, fix_reshape=False):
    clipped_block_size = min(block_size, H, W)

    if fix_reshape:
        # This should match the .reshape() dimension order below.
        h_i, w_i = ndgrid(torch.arange(H), torch.arange(W))
    else:
        # The original produces crazy stride patterns, due to .reshape() offset winding.
        # This is only visible when H != W.
        w_i, h_i = ndgrid(torch.arange(W), torch.arange(H))

    valid_block = ((w_i >= clipped_block_size // 2) & (w_i < W - (clipped_block_size - 1) // 2)) & \
                 ((h_i >= clipped_block_size // 2) & (h_i < H - (clipped_block_size - 1) // 2))

    valid_block = torch.reshape(valid_block, (1, 1, H, W))

    return valid_block

def main():
    common_args = dict(H=4, W=5, block_size=3)

    for fix in [False, True]:
        args = dict(H=4, W=5, block_size=3, fix_reshape=fix)
        grid = valid_block(**args)
        print(args)
        print(f"{grid.shape=}")
        print(grid)
        print()

if __name__ == "__main__":
    main()
```

Author: crutcher

tests/layers/__init__.py

@@ -0,0 +1,47 @@
+import importlib
+import os
+import unittest
+
+import torch
+
+from timm.layers import drop
+
+torch_backend = os.environ.get('TORCH_BACKEND')
+if torch_backend is not None:
+    importlib.import_module(torch_backend)
+torch_device = os.environ.get('TORCH_DEVICE', 'cpu')
+
+class ClipMaskTests(unittest.TestCase):
+    def test_clip_mask_2d_odd(self):
+        mask = drop.clip_mask_2d(h=5, w=7, kernel=3, device=torch_device)
+        assert mask.device == torch.device(torch_device)
+        assert mask.tolist() == \
+            [
+                [False, False, False, False, False, False, False],
+                [False, True, True, True, True, True, False],
+                [False, True, True, True, True, True, False],
+                [False, True, True, True, True, True, False],
+                [False, False, False, False, False, False, False],
+            ]
+
+    def test_clip_mask_2d_even(self):
+        mask = drop.clip_mask_2d(h=5, w=7, kernel=2, device=torch_device)
+        assert mask.device == torch.device(torch_device)
+        # TODO: This is a suprising result; should even kernels be forbidden?
+        assert mask.tolist() == \
+            [
+                [False, False, False, False, False, False, False],
+                [False, True, True, True, True, True, True],
+                [False, True, True, True, True, True, True],
+                [False, True, True, True, True, True, True],
+                [False, True, True, True, True, True, True],
+            ]
+
+    def test_clip_mask_2d_kernel_too_big(self):
+        try:
+            drop.clip_mask_2d(h=4, w=7, kernel=5, device=torch_device)
+            raise RuntimeError("Expected throw")
+
+        except AssertionError as e:
+            assert "kernel=5 > min(h=4, w=7)" in e.args[0]
+

tests/layers/test_drop.py

@@ -21,6 +21,43 @@
 from .grid import ndgrid
 
 
+def clip_mask_2d(
+        h: int,
+        w: int,
+        kernel: int,
+        device,
+):
+    """Build a clip mask.
+
+    Returns a mask of all points which permit a (kernel, kernel) sized
+    block to sit entirely within the (h, w) index space.
+
+    Requires `kernel <= min(h, w)`.
+
+    TODO: Should even kernels be forbidden?
+    Even kernels behave oddly, but are not forbidden for historical reasons.
+
+    Args:
+        h: the height.
+        w: the width.
+        kernel_size: the size of the kernel.
+        device: the target device.
+        check_kernel: when true, assert that the kernel_size is odd.
+
+    Returns:
+        a (h, w) bool mask tensor.
+    """
+    assert kernel <= min(h, w), f"{kernel=} > min({h=}, {w=})"
+
+    h_i, w_i = ndgrid(torch.arange(h, device=device), torch.arange(w, device=device))
+    return (
+        (h_i >= kernel // 2) &
+        (h_i < h - (kernel - 1) // 2) &
+        (w_i >= kernel // 2) &
+        (w_i < w - (kernel - 1) // 2)
+    ).reshape(h, w)
+
+
 def drop_block_2d(
         x,
         drop_prob: float = 0.1,
@@ -30,23 +67,39 @@ def drop_block_2d(
         inplace: bool = False,
         batchwise: bool = False
 ):
-    """ DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
+    """DropBlock. See https://arxiv.org/pdf/1810.12890.pdf
 
     DropBlock with an experimental gaussian noise option. This layer has been tested on a few training
     runs with success, but needs further validation and possibly optimization for lower runtime impact.
+
+    Args:
+        drop_prob: the probability of dropping any given block.
+        block_size: the size of the dropped blocks; should be odd.
+        gamma_scale: adjustment scale for the drop_prob.
+        with_noise: should normal noise be added to the dropped region?
+        inplace: if the drop should be applied in-place on the input tensor.
+        batchwise: should the entire batch use the same drop mask?
+
+    Returns:
+        If inplace, the modified `x`; otherwise, the dropped copy of `x`, on the same device.
     """
     B, C, H, W = x.shape
     total_size = W * H
-    clipped_block_size = min(block_size, min(W, H))
+
+    # TODO: This behaves oddly when clipped_block_size < block_size, or block_size % 2 == 0.
+    clipped_block_size = min(block_size, W, H)
+
     # seed_drop_rate, the gamma parameter
     gamma = gamma_scale * drop_prob * total_size / clipped_block_size ** 2 / (
             (W - block_size + 1) * (H - block_size + 1))
 
     # Forces the block to be inside the feature map.
-    w_i, h_i = ndgrid(torch.arange(W, device=x.device), torch.arange(H, device=x.device))
-    valid_block = ((w_i >= clipped_block_size // 2) & (w_i < W - (clipped_block_size - 1) // 2)) & \
-                  ((h_i >= clipped_block_size // 2) & (h_i < H - (clipped_block_size - 1) // 2))
-    valid_block = torch.reshape(valid_block, (1, 1, H, W)).to(dtype=x.dtype)
+    valid_block = clip_mask_2d(
+        h=H,
+        w=W,
+        kernel=clipped_block_size,
+        device=x.device,
+    ).reshape((1, 1, H, W)).to(dtype=x.dtype)
 
     if batchwise:
         # one mask for whole batch, quite a bit faster