RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Test MiniMax output attention patches

This commit is contained in:
Pyro 2026-08-04 01:41:02 +02:00
parent 710b6b7ee0
commit 13de02c097
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tests-unit/comfy_test/test_minimax_attention.py
View File

@ -26,6 +26,10 @@ def test_attention_patch_accepts_tuple_and_mapping_callbacks(monkeypatch):
assert extra_options["n_heads"] == 2
return {"q": q * 2, "v": v * 3}
def output_patch(out, extra_options):
assert extra_options["block_index"] == 3
return out + 4
def fake_attention(q, k, v, *args, **kwargs):
seen.update(q=q, k=k, v=v)
return v.transpose(1, 2).reshape(1, v.shape[2], -1)
@ -36,11 +40,14 @@ def test_attention_patch_accepts_tuple_and_mapping_callbacks(monkeypatch):
x,
transformer_options={
"block_index": 3,
"patches": {"attn1_patch": [tuple_patch, mapping_patch]},
"patches": {
"attn1_patch": [tuple_patch, mapping_patch],
"attn1_output_patch": [output_patch],
},
},
)
assert torch.equal(seen["q"], torch.full((1, 2, 2, 2), 2.0))
assert torch.equal(seen["k"], torch.full((1, 2, 2, 2), 2.0))
assert torch.equal(seen["v"], torch.full((1, 2, 2, 2), 9.0))
assert torch.equal(output, torch.full((2, 4), 9.0))
assert torch.equal(output, torch.full((2, 4), 13.0))