Repository: facebookresearch/xformers

Issue:
[CI] InProj countainer not properly covered
# 🐛 Bug
Not a bug per say, but many inProjCountainer options are not covered, see https://app.codecov.io/gh/facebookresearch/xformers/blob/main/xformers/components/in_proj_container.py



Discussion:
~~hmm, looks like it's actually tested but it's yet another space that the code coverage misses (the backward implementations are called by autograd from outside the original python context..)~~
closed by mistake, comment applies to the other related issue (reversible BW)

Resolving patch:
diff --git a/xformers/components/__init__.py b/xformers/components/__init__.py
--- a/xformers/components/__init__.py
+++ b/xformers/components/__init__.py
@@ -12,6 +12,7 @@
 
 from .activations import Activation, build_activation  # noqa
 from .attention import Attention, build_attention  # noqa
+from .in_proj_container import InProjContainer, InProjParams  # noqa
 from .multi_head_dispatch import MultiHeadDispatch, MultiHeadDispatchConfig  # noqa
 from .residual import LayerNormStyle, PostNorm, PreNorm, Residual  # noqa
 


Test patch:
diff --git a/tests/test_attentions.py b/tests/test_attentions.py
--- a/tests/test_attentions.py
+++ b/tests/test_attentions.py
@@ -6,7 +6,7 @@
 import pytest
 import torch
 
-from xformers.components import MultiHeadDispatch
+from xformers.components import InProjContainer, InProjParams, MultiHeadDispatch
 
 # Automatically test all the registered attentions
 from xformers.components.attention import (
@@ -164,6 +164,71 @@ def test_kqv_ordering(
     assert torch.allclose(res_false[0, :, :], res_false[1, :, :])
 
 
+@pytest.mark.parametrize("small_init", [False, True])
+@pytest.mark.parametrize("proj_bias", [False, True])
+@pytest.mark.parametrize("same_sizes", [False, True])
+@pytest.mark.parametrize("same_settings", [False, True])
+def test_inproj(
+    small_init: bool, proj_bias: bool, same_sizes: bool, same_settings: bool
+):
+
+    test_config = {
+        "name": "scaled_dot_product",
+        "dropout": 0.1,
+        "causal": False,
+        "seq_len": SEQ,
+        "window_size": SEQ // 8 + 1,
+        "num_heads": 1,
+        "dim_head": MODEL,
+    }
+
+    attention = build_attention(test_config)
+
+    # Construct the initial projection, test different options
+    in_params = InProjParams(MODEL, MODEL, proj_bias, small_init)
+
+    if same_settings:
+        in_proj = InProjContainer(in_params, None, None)
+    else:
+        out_features = MODEL if same_sizes else MODEL - 16
+        in_params_flip = InProjParams(MODEL, out_features, not proj_bias, small_init)
+        in_proj = InProjContainer(in_params, in_params_flip, in_params_flip)
+
+    # build a multi head dispatch to test this attention mechanism
+    multi_head = MultiHeadDispatch(
+        seq_len=SEQ,
+        dim_model=MODEL,
+        residual_dropout=0.1,
+        num_heads=1,
+        attention=attention,
+        in_proj_container=in_proj,
+    )
+
+    # Check kqv are not flipped
+    # this will not catch all issues, but would catch a V being misplaced
+    # make k and q complimentary, so that QKt is all zero and attention is uniform
+
+    q = torch.cat(
+        (
+            torch.rand((1, MODEL // 2)),
+            torch.zeros((1, MODEL // 2)),
+        ),
+        dim=1,
+    ).expand((BATCH, SEQ, MODEL))
+
+    k = torch.cat(
+        (
+            torch.zeros((1, MODEL // 2)),
+            torch.rand((1, MODEL // 2)),
+        ),
+        dim=1,
+    ).expand((BATCH, SEQ, MODEL))
+    v = torch.rand(BATCH, SEQ, MODEL)
+
+    # just check that a FW does not assert out
+    _ = multi_head(query=q, key=k, value=v)
+
+
 @pytest.mark.parametrize("heads", [1, 4])
 @pytest.mark.parametrize("attention_name", ATTENTION_REGISTRY.keys())
 @pytest.mark.parametrize("device", DEVICES)