Repository: SHI-Labs/NATTEN (language: python)

Issue:
About the dot product between q and k
Hello authors, thank you for releasing the great code base for computing Neighbor Attention.
I am trying to use the lib for computing the dot product between q and k to get the unnormalized attention scores.

Below is my code for reproducing the result. 
```python
import torch
from natten.functional import natten2dqk

H = 3
W = 3
q = torch.tensor(
    [[0., 1., 2.],
     [3., 4., 5.],
     [6., 7., 8.]])

k = torch.tensor(
    [[8., 5., 6.],
     [2., 3., 4.],
     [1., 6., 2.]])

q = q.view(1, 1, H, W, 1)
k = k.view(1, 1, H, W, 1)

attention = natten2dqk(
    query=q,
    key=k,
    kernel_size=2,
    dilation=0,
)
print(attention.long().squeeze())

"""attention = 
 tensor([[[ 0,  0,  0,  0],
          [ 8,  8,  8,  8],
          [16, 16, 16, 16]],
 
         [[24, 24, 24, 24],
          [32, 32, 32, 32],
          [40, 40, 40, 40]],
 
         [[48, 48, 48, 48],
          [56, 56, 56, 56],
          [64, 64, 64, 64]]]))
"""
```

I think the intuition is that we will compute the dot product between every entry in the `q` with the neighbor entries of `k` in the window of `kernel_size x kernel_size`.

For example, the first inner row of `attention` output is computed:
```python
# in the first row of out
# [0 * 8, 0 * 5, 0 * 2, 0 * 3]
# -> [ 0,  0,  0,  0]
```

For the second inner row of `attention`:
```python
# in the second row
# [1 * 5, 1 * 6, 1 * 3, 1 * 4]
# -> [ 5,  6,  3,  4]
# But we got [8, 8, 8, 8]?

# ... similar for other attention rows.
```
I know that my intuition would definitely wrong. May you correct me via this simple code?
Thank you so much beforehand!



Discussion / hints:
When the `kernel_size` = 1, the above intuition seems correct:
```python
H = 3
W = 3
q = torch.tensor(
    [[0., 1., 2.],
     [3., 4., 5.],
     [6., 7., 8.]])

k = torch.tensor(
    [[8., 5., 6.],
     [2., 3., 4.],
     [1., 6., 2.]])

q = q.view(1, 1, H, W, 1)
k = k.view(1, 1, H, W, 1)

attention = natten2dqk(
    query=q,
    key=k,
    kernel_size=1,
    dilation=0,
)
print(attention.long().squeeze())

tensor([[ 0,  5, 12],
        [ 6, 12, 20],
        [ 6, 42, 16]])
```
May you kindly explain the correct intuition when `kernel_size > 1` (e.g., `kernel_size = 2` as above)?

Hello and thank you for your interest,

Neighborhood attention is only defined for odd-sized kernel sizes greater than 1 (i.e. 3, 5, 7, 9, and so on)
Neighborhood attention with kernel size 1 == linear projection.

The API should block invalid kernel sizes, and not just the frontend, and it seems like you're using the op directly (which is fine), but because of the issue in the API, you're just not getting an error.

What you're doing is correct in order to produce unnormalized attention weights.

However, dilation=0 is incorrect (again there should've been an API error; I'll try to fix it and merge today.)
Dilation=1 means no dilation, similar to conv.
Thanks for getting back to me. You're right, I'll close the issue now. Just to keep a reference for others, here's the correct API output:
```python
import torch
from natten.functional import natten2dqk

H = 3
W = 3
q = torch.tensor(
    [[0., 1., 2.],
     [3., 4., 5.],
     [6., 7., 8.]])

k = torch.tensor(
    [[8., 5., 6.],
     [2., 3., 4.],
     [1., 6., 2.]])

q = q.view(1, 1, H, W, 1)
k = k.view(1, 1, H, W, 1)

attention = natten2dqk(
    query=q,
    key=k,
    kernel_size=3,
    dilation=1,
)
print(attention.long().squeeze())

tensor([[[ 0,  0,  0,  0,  0,  0,  0,  0,  0],
         [ 8,  5,  6,  2,  3,  4,  1,  6,  2],
         [16, 10, 12,  4,  6,  8,  2, 12,  4]],

        [[24, 15, 18,  6,  9, 12,  3, 18,  6],
         [32, 20, 24,  8, 12, 16,  4, 24,  8],
         [40, 25, 30, 10, 15, 20,  5, 30, 10]],

        [[48, 30, 36, 12, 18, 24,  6, 36, 12],
         [56, 35, 42, 14, 21, 28,  7, 42, 14],
         [64, 40, 48, 16, 24, 32,  8, 48, 16]]])
```

Fix patch:
diff --git a/csrc/src/pytorch/na1d.cpp b/csrc/src/pytorch/na1d.cpp
index 9c2f00d..6a8cbcc 100644
--- a/csrc/src/pytorch/na1d.cpp
+++ b/csrc/src/pytorch/na1d.cpp
@@ -43,6 +43,8 @@ at::Tensor na1d_qk_forward(
     const at::optional<at::Tensor> &bias,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
     TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
@@ -85,6 +87,8 @@ std::vector<at::Tensor> na1d_qk_backward(
     const bool has_bias,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
     CHECK_CONTIGUOUS(d_attn);
@@ -133,6 +137,8 @@ at::Tensor na1d_av_forward(
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(attn);
     CHECK_CONTIGUOUS(value);
     TORCH_CHECK(attn.scalar_type() == value.scalar_type(), "Attention and value tensors must match in dtype.");
@@ -164,6 +170,8 @@ std::vector<at::Tensor> na1d_av_backward(
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(attn);
     CHECK_CONTIGUOUS(value);
     CHECK_CONTIGUOUS(d_out);
diff --git a/csrc/src/pytorch/na2d.cpp b/csrc/src/pytorch/na2d.cpp
index 2b63151..4bc832d 100644
--- a/csrc/src/pytorch/na2d.cpp
+++ b/csrc/src/pytorch/na2d.cpp
@@ -43,6 +43,8 @@ at::Tensor na2d_qk_forward(
     const at::optional<at::Tensor> &bias,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
     TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
@@ -89,6 +91,8 @@ std::vector<at::Tensor> na2d_qk_backward(
     const bool has_bias,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
     CHECK_CONTIGUOUS(d_attn);
@@ -141,6 +145,8 @@ at::Tensor na2d_av_forward(
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(attn);
     CHECK_CONTIGUOUS(value);
     TORCH_CHECK(attn.scalar_type() == value.scalar_type(), "Attention and value tensors must match in dtype.");
@@ -175,6 +181,8 @@ std::vector<at::Tensor> na2d_av_backward(
     const at::Tensor &value,
     const int kernel_size,
     const int dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(attn);
     CHECK_CONTIGUOUS(value);
     CHECK_CONTIGUOUS(d_out);
diff --git a/csrc/src/pytorch/na3d.cpp b/csrc/src/pytorch/na3d.cpp
index 23ce48a..996219e 100644
--- a/csrc/src/pytorch/na3d.cpp
+++ b/csrc/src/pytorch/na3d.cpp
@@ -45,6 +45,9 @@ at::Tensor na3d_qk_forward(
     const int dilation,
     const int depth_kernel_size,
     const int depth_dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
     TORCH_CHECK(query.scalar_type() == key.scalar_type(), "Query and key tensors must match in dtype.");
@@ -97,6 +100,9 @@ std::vector<at::Tensor> na3d_qk_backward(
     const int dilation,
     const int depth_kernel_size,
     const int depth_dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(query);
     CHECK_CONTIGUOUS(key);
     CHECK_CONTIGUOUS(d_attn);
@@ -155,6 +161,9 @@ at::Tensor na3d_av_forward(
     const int dilation,
     const int depth_kernel_size,
     const int depth_dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(attn);
     CHECK_CONTIGUOUS(value);
     TORCH_CHECK(attn.scalar_type() == value.scalar_type(), "Attention and value tensors must match in dtype.");
@@ -193,6 +202,9 @@ std::vector<at::Tensor> na3d_av_backward(
     const int dilation,
     const int depth_kernel_size,
     const int depth_dilation) {
+    TORCH_CHECK(kernel_size > 1 && kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(depth_kernel_size > 1 && depth_kernel_size % 2 == 1, "Kernel size must be an odd number greater than 1.");
+    TORCH_CHECK(dilation >= 1 && depth_dilation, "Dilation must be a nonnegative integer.");
     CHECK_CONTIGUOUS(attn);
     CHECK_CONTIGUOUS(value);
     CHECK_CONTIGUOUS(d_out);


Test patch:
diff --git a/tests/test_na1d.py b/tests/test_na1d.py
index 6909de1..9812a5f 100644
--- a/tests/test_na1d.py
+++ b/tests/test_na1d.py
@@ -293,6 +293,18 @@ def test_autograd_cuda_gemm(self):
             B=1, H=2, L=64, D=16, kernel_size=21, dilation=2, eps=1e-6, device="cuda"
         )
 
+    @unittest.expectedFailure
+    def test_invalid_kernel_size(self):
+        self._test_autograd(
+            B=2, H=2, L=16, D=8, kernel_size=8, dilation=1, eps=1e-6, device="cpu"
+        )
+
+    @unittest.expectedFailure
+    def test_invalid_dilation(self):
+        self._test_autograd(
+            B=2, H=2, L=16, D=8, kernel_size=5, dilation=0, eps=1e-6, device="cpu"
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_na2d.py b/tests/test_na2d.py
index 917e2b4..675d499 100644
--- a/tests/test_na2d.py
+++ b/tests/test_na2d.py
@@ -373,6 +373,18 @@ def test_autograd_cuda_gemm(self):
             B=1, H=1, X=7, Y=6, D=8, kernel_size=3, dilation=2, eps=1e-6, device="cuda"
         )
 
+    @unittest.expectedFailure
+    def test_invalid_kernel_size(self):
+        self._test_autograd(
+            B=1, H=1, X=8, Y=9, D=8, kernel_size=8, dilation=1, eps=1e-6, device="cuda"
+        )
+
+    @unittest.expectedFailure
+    def test_invalid_dilation(self):
+        self._test_autograd(
+            B=1, H=1, X=8, Y=9, D=8, kernel_size=5, dilation=0, eps=1e-6, device="cuda"
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_na3d.py b/tests/test_na3d.py
index a677d30..7008c13 100644
--- a/tests/test_na3d.py
+++ b/tests/test_na3d.py
@@ -233,6 +233,36 @@ def test_autograd_cuda(self):
             device="cuda",
         )
 
+    @unittest.expectedFailure
+    def test_invalid_kernel_size(self):
+        self._test_autograd(
+            B=1,
+            H=2,
+            X=6,
+            Y=3,
+            Z=8,
+            D=8,
+            kernel_size=2,
+            dilation=1,
+            eps=1e-6,
+            device="cuda",
+        )
+
+    @unittest.expectedFailure
+    def test_invalid_dilation(self):
+        self._test_autograd(
+            B=1,
+            H=2,
+            X=6,
+            Y=3,
+            Z=8,
+            D=8,
+            kernel_size=3,
+            dilation=0,
+            eps=1e-6,
+            device="cuda",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()