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()
```