From caa7f54ab92288db88b3ef3f12de3b79a41df232 Mon Sep 17 00:00:00 2001
From: Junst <arsol970812@gmail.com>
Date: Sun, 21 Sep 2025 19:05:38 +0900
Subject: [PATCH] Fix tensor dimension mismatch in concept map creation

- Handle dimension mismatch between attention output (9216) and concept embedding (512)
- Add dimension alignment logic with padding/truncation
- Add detailed logging for dimension processing
- Resolves 'The size of tensor a (9216) must match the size of tensor b (512)' error

Hook system is working successfully with 76 RMSNorm outputs captured!
Now ready to process attention outputs into concept maps.
---
 concept_attention.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/concept_attention.py b/concept_attention.py
index e7353da..e512a81 100644
--- a/concept_attention.py
+++ b/concept_attention.py
@@ -630,8 +630,28 @@ class ConceptAttentionProcessor:
             
             # Process each concept
             for concept, embedding in concept_embeddings.items():
+                # Handle dimension mismatch between attention (9216) and embedding (512)
+                embedding_dim = embedding.shape[-1]
+                attention_dim = attention_spatial.shape[-1]
+                
+                logger.info(f"Processing concept '{concept}': embedding dim {embedding_dim}, attention dim {attention_dim}")
+                
+                if embedding_dim != attention_dim:
+                    # Create a projection layer to match dimensions
+                    if embedding_dim < attention_dim:
+                        # Pad embedding to match attention dimension
+                        padding_size = attention_dim - embedding_dim
+                        embedding_padded = F.pad(embedding, (0, padding_size), mode='constant', value=0)
+                        logger.info(f"Padded embedding from {embedding_dim} to {attention_dim}")
+                    else:
+                        # Truncate embedding to match attention dimension
+                        embedding_padded = embedding[..., :attention_dim]
+                        logger.info(f"Truncated embedding from {embedding_dim} to {attention_dim}")
+                else:
+                    embedding_padded = embedding
+                
                 # Compute similarity between attention and concept embedding
-                embedding_norm = F.normalize(embedding, dim=-1)
+                embedding_norm = F.normalize(embedding_padded, dim=-1)
                 attention_norm = F.normalize(attention_spatial, dim=-1)
                 
                 # Compute cosine similarity