"""Contrastive pre-training of an unitary encoder"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../src/models/clip/unitary_clip.ipynb.

# %% auto #0
__all__ = ['RotaryMultiheadAttention', 'FeedForwardBlock', 'UnitaryEncoderAttnBlock', 'UnitaryTextEncoderConfig',
           'UnitaryTextEncoder', 'SelfAttnBlock', 'PackingTransformer', 'CoreTransformer', 'CircuitEncoderConfig',
           'CircuitEncoder', 'UnitaryCLIPConfig', 'UnitaryCLIP']

# %% ../../../src/models/clip/unitary_clip.ipynb #06272f6f-b4e3-4504-a90a-feebbf6ad821
from ...imports import *
from ..config_model import ConfigModel
import genQC.models.transformers.attention as attn
import genQC.models.layers as layers
from ..position_encoding import LearnedPositionalEmbedding, RotaryPositionalEmbedding, RotaryPositionalEmbedding2D

# %% ../../../src/models/clip/unitary_clip.ipynb #909ad61a-11c8-4e5f-a4d3-7dd80fce577a
class RotaryMultiheadAttention(nn.Module):
    """
    MultiheadAttention described in the paper: Attention Is All You Need (https://arxiv.org/abs/1706.03762).
    We add a rotary position encoding (RoPE). 

    The attention core is `F.scaled_dot_attention` from pytorch. 
    Could be switched to `https://github.com/Dao-AILab/flash-attention` or `xFormers`.
    """

    def __init__(self,
                 in_dim: int,
                 embed_dim: int, 
                 num_heads: int, 
                 bias: bool = True, 
                 p_rope: float = 1.0, 
                 max_seq_len: int = 4096, 
                 base_rope: float = 10_000,
                 enable_qk_norm: bool = False) -> None:
        
        super().__init__()

        self.num_heads = num_heads
        self.bias      = bias
        self.head_dim  = embed_dim // num_heads 

        self.q_proj   = nn.Linear(in_dim, embed_dim, bias=bias)
        self.k_proj   = nn.Linear(in_dim, embed_dim, bias=bias)
        self.v_proj   = nn.Linear(in_dim, embed_dim, bias=bias)
        
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        self.enable_qk_norm = enable_qk_norm
        if self.enable_qk_norm:
            self.q_norm = nn.RMSNorm(self.head_dim)
            self.k_norm = nn.RMSNorm(self.head_dim)
        
        self.rope = RotaryPositionalEmbedding(head_dim=self.head_dim, p=p_rope, max_seq_len=max_seq_len, base=base_rope)
        
        self._init_weights()
          
    def _init_weights(self) -> None:
        nn.init.xavier_normal_(self.q_proj.weight)
        nn.init.xavier_normal_(self.k_proj.weight)
        nn.init.xavier_normal_(self.v_proj.weight)
        nn.init.xavier_normal_(self.out_proj.weight)

        if self.bias:
            nn.init.zeros_(self.q_proj.bias)
            nn.init.zeros_(self.k_proj.bias)
            nn.init.zeros_(self.v_proj.bias)
            nn.init.zeros_(self.out_proj.bias)

    
    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, pos_idx: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Assumes batch first. When `pos_idx` is provided we use RoPE, else NOT!

        Shapes:
            query     ... [b, n1, c]
            key/value ... [b, n2, c]
        """

        assert key.shape == value.shape
        
        b, n1, _ = query.shape
        _, n2, _ = key.shape

        q = self.q_proj(query)
        k = self.k_proj(key)
        v = self.v_proj(value)

        q = q.view(b, n1, self.num_heads, self.head_dim)
        k = k.view(b, n2, self.num_heads, self.head_dim)
        v = v.view(b, n2, self.num_heads, self.head_dim)

        if self.enable_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        
        if exists(pos_idx):
            q = self.rope(q, pos_idx=pos_idx)
            k = self.rope(k, pos_idx=pos_idx)

        # scaled_dot_product_attention takes [b, num_heads, seq, head_dim]
        q = q.permute((0, 2, 1, 3)) 
        k = k.permute((0, 2, 1, 3)) 
        v = v.permute((0, 2, 1, 3)) 
        
        # see https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
        attn = F.scaled_dot_product_attention(query=q, 
                                              key=k, 
                                              value=v, 
                                              attn_mask=None, 
                                              dropout_p=0.0,
                                              is_causal=False, 
                                              scale=None, 
                                              #enable_gqa=False
                                             )

        # back to [b, seq, num_heads, head_dim]
        attn = attn.permute((0, 2, 1, 3)) 

        # pack heads together
        attn = attn.reshape(b, n1, self.num_heads * self.head_dim)
        attn = self.out_proj(attn)
        return attn

# %% ../../../src/models/clip/unitary_clip.ipynb #a6bd03c7-70e4-4264-8a08-b5ee2a294f8c
class FeedForwardBlock(nn.Module):
    """
    A small dense feed-forward network as used in `transformers`. Assumes channel last.
    Inspired by https://arxiv.org/pdf/2401.11605 and added 
    from https://arxiv.org/pdf/2002.05202 a modification to SiGLU structure.
    """

    def __init__(self, in_dim: int, hidden_dim: int, dropout: float = 0.0) -> None:
        super().__init__()
        self.hidden_dim = hidden_dim
        self.proj_in  = nn.Linear(in_dim, 2*hidden_dim) # factor two for GLU part split
        self.proj_out = nn.Linear(hidden_dim, in_dim) 
        self.act   = nn.SiLU()
        self.drop  = nn.Dropout(dropout)

        self._init_weights()
   
    def _init_weights(self) -> None:
        nn.init.zeros_(self.proj_out.bias)
        # nn.init.xavier_normal_(self.proj_out.weight)
    
    def siglu(self, x: torch.Tensor) -> torch.Tensor:
        x = self.proj_in(x) 
        return x[..., :self.hidden_dim] * self.act(x[..., self.hidden_dim:])

    #@torch.compile
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.siglu(x)
        x = self.drop(x)
        x = self.proj_out(x)
        return x

# %% ../../../src/models/clip/unitary_clip.ipynb #0b4b123f-40cd-4321-ab5e-9427dd907396
class UnitaryEncoderAttnBlock(nn.Module):
    """A self-attention block with 2d-RoPE."""
    
    def __init__(self, 
                 ch: int, 
                 y_emb_size: int,
                 num_heads: int, 
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()

        self.self_att = RotaryMultiheadAttention(in_dim=ch+y_emb_size, embed_dim=ch, num_heads=num_heads, p_rope=p_rope, base_rope=base_rope)
        self.self_att.rope = RotaryPositionalEmbedding2D(head_dim=self.self_att.head_dim, p=p_rope, base=base_rope)
        
        self.ff        = FeedForwardBlock(in_dim=ch, hidden_dim=2*ch)   
        self.norm_self = nn.RMSNorm(ch)
        self.norm_ff   = nn.RMSNorm(ch)
        self.drop      = nn.Dropout(dropout)
        
        self._init_weights()
          
    def _init_weights(self) -> None:

        # note a bonus of res-pos-norm is that we can init as identity!
        nn.init.zeros_(self.norm_self.weight)  
        nn.init.zeros_(self.norm_ff.weight)

    def forward(self, x: torch.Tensor, y_emb: torch.Tensor, pos_idx: torch.Tensor) -> torch.Tensor:
        """
        Assumes batch first.
        
        Shapes:
            x       ... [b, n, ch1]       
            y_emb   ... [b, n, ch2]
            pos_idx ... [b, n, 2] or [n, 2]
        """

        # Self-attention part
        self_out = torch.cat([x, y_emb], dim=-1)
        self_out = self.self_att(query=self_out, key=self_out, value=self_out, pos_idx=pos_idx)
        self_out = self.norm_self(self_out)  
        self_out = self.drop(self_out) + x      

        # Feed-Forward part
        feed_out = self.ff(self_out)
        feed_out = self.norm_ff(feed_out) 
        feed_out = self.drop(feed_out) + self_out                      
        return feed_out

# %% ../../../src/models/clip/unitary_clip.ipynb #b68db894-85a3-4d58-b6c3-1f44075f7e7d
@dataclass
class UnitaryTextEncoderConfig: 
    text_embed_ch: int
    text_encoding_ch: int  
    text_attn_num_heads: int  
    text_attn_depth: int
    
    unitary_encoding_ch: int  
    unitary_downscale_factor: int
    
    main_num_heads: int
    main_depth: int
    
    use_rope: bool
    p_rope: float
    base_rope: float
    dropout: float

# %% ../../../src/models/clip/unitary_clip.ipynb #f61c338b-4f43-4d18-a2b2-0cba26bad9a4
class UnitaryTextEncoder(ConfigModel):
    def __init__(self,     
                 text_embed_ch: int,
                 text_encoding_ch: int,
                 text_attn_num_heads: int,
                 text_attn_depth: int,
                 unitary_encoding_ch: int,
                 unitary_downscale_factor: int,
                 main_num_heads: int,
                 main_depth: int,
                 use_rope: bool,
                 p_rope: float,
                 base_rope: float,
                 dropout: float) -> None:  
        """
        text_embed_ch ... number of channels of the input text encodings `y_emb`

        The text channels `text_encoding_ch` are concatenated with the unitary channels `unitary_encoding_ch`.
        """
        super().__init__()   

        self.params_config = UnitaryTextEncoderConfig(text_embed_ch=text_embed_ch,
                                                      text_encoding_ch=text_encoding_ch,
                                                      text_attn_num_heads=text_attn_num_heads,
                                                      text_attn_depth=text_attn_depth,
                                                      unitary_encoding_ch=unitary_encoding_ch,
                                                      unitary_downscale_factor=unitary_downscale_factor,
                                                      main_num_heads=main_num_heads,
                                                      main_depth=main_depth,
                                                      use_rope=use_rope,
                                                      p_rope=p_rope,
                                                      base_rope=base_rope,
                                                      dropout=dropout)
        
        # Text pre-process
        self.text_proj        = nn.Linear(text_embed_ch, text_encoding_ch)
        self.text_norm        = nn.RMSNorm(text_encoding_ch)
        
        self.text_attn_blocks = nn.ModuleList([attn.BasisSelfAttnBlock(ch=text_encoding_ch,
                                                                       num_heads=text_attn_num_heads, 
                                                                       dropout=dropout,
                                                                       batch_first=True) 
                                               for d in range(text_attn_depth)
                                              ])
        
        # Unitary pre-process
        self.unitary_proj      = nn.Conv2d(2, unitary_encoding_ch, kernel_size=1, stride=1, padding="same") 
        self.unitary_downscale = nn.PixelUnshuffle(unitary_downscale_factor)
        self.unitary_downscale_factor = unitary_downscale_factor

        self.use_rope = use_rope
        if not self.use_rope:
            self.unitary_pos_enc = layers.PositionalEncoding2D(d_model=unitary_encoding_ch, freq_factor=1_000) 
        
        # Main transformer
        self.encoding_ch = unitary_encoding_ch * (unitary_downscale_factor**2)
        
        self.transformer_blocks = nn.ModuleList([UnitaryEncoderAttnBlock(ch=self.encoding_ch,
                                                                         y_emb_size=text_encoding_ch,
                                                                         num_heads=main_num_heads, 
                                                                         dropout=dropout,
                                                                         p_rope=p_rope,
                                                                         base_rope=base_rope)         
                                                 for d in range(main_depth)
                                                ])

        self.norm_final = nn.RMSNorm(self.encoding_ch)
        
        print(f"[INFO]: Creating `UnitaryTextEncoder` with `{unitary_downscale_factor=}` and `encoding_ch={self.encoding_ch}`.")
        self._init_weights()
          
    def _init_weights(self) -> None:
        # nn.init.xavier_normal_(self.text_proj.weight)
        # nn.init.xavier_normal_(self.unitary_proj.weight)

        nn.init.zeros_(self.text_proj.bias)
        nn.init.zeros_(self.unitary_proj.bias)    

    def preproc_text(self, y_emb):
        y_emb = self.text_proj(y_emb)   # ... [batch, seq_y, text_encoding_ch]

        for text_attn_block in self.text_attn_blocks:
            y_emb = text_attn_block(y_emb)

        return y_emb
    
    def preproc_unitary(self, U):
        u_emb = self.unitary_proj(U)             # ... [batch, unitary_encoding_ch, N, N]
        if not self.use_rope:
            u_emb = self.unitary_pos_enc(u_emb)
        u_emb = self.unitary_downscale(u_emb)    # ... [batch, unitary_encoding_ch * r^2, N/r, N/r]

        # Reshape and permute from image to sentence shape
        b, ch, *_ = u_emb.shape
        u_emb = torch.reshape(u_emb, (b, ch, -1))    #  to [batch, unitary_encoding_ch * r^2, (N/r)^2]
        u_emb = torch.permute(u_emb, (0, 2, 1))      #  to [batch, (N/r)^2, unitary_encoding_ch * r^2]

        return u_emb

    def forward(self, y_emb: torch.Tensor, U: torch.Tensor, pool: bool = False, penultimate: bool = False) -> torch.Tensor: 
        """
        penultimate_output = False ... take all attn layers
        penultimate_output = True  ... skip the last attn layers
        
        Shapes:
            y_emb  ... [b, seq, text_embed_ch]  
            U      ... [b, 2, N, N]
        """
        
        # Pre-process multimodial inputs
        x     = self.preproc_unitary(U)    # ... [batch, seq_u, unitary_encoding_ch * r^2]
        y_emb = self.preproc_text(y_emb)   # ... [batch, seq_y, text_encoding_ch]
        
        y_emb = y_emb.mean(dim=1, keepdim=True)            # ... [batch,     1, text_encoding_ch]
        y_emb = self.text_norm(y_emb)
        y_emb = y_emb.expand(x.shape[0], x.shape[1], -1)   # ... [batch, seq_u, text_encoding_ch]
        
        # Main transformer pass
        if self.use_rope:
            N   = U.shape[-1] // self.unitary_downscale_factor
            pos = torch.arange(N).expand(N, -1)
            pos_idx = torch.stack([pos.T, pos], dim=-1).reshape(-1, 2)  # ... [seq_u, 2]
        else:
            pos_idx = None
        
        if not penultimate:
            for transformer_block in self.transformer_blocks:
                x = transformer_block(x, y_emb=y_emb, pos_idx=pos_idx)
                
        else:
            for transformer_block in self.transformer_blocks[:-1]:
                x = transformer_block(x, y_emb=y_emb, pos_idx=pos_idx)
        
        if pool: 
            x = torch.mean(x, dim=1) # [batch, ch]   

        x = self.norm_final(x)
        return x

# %% ../../../src/models/clip/unitary_clip.ipynb #5d326b93-4ed2-48f2-982c-80471a428751
class SelfAttnBlock(nn.Module):
    """A self-attention block with RoPE."""
    
    def __init__(self, ch: int, num_heads: int, dropout: float = 0.0, p_rope: float = 1.0, base_rope: float = 10_000) -> None:
        super().__init__()

        self.self_att = RotaryMultiheadAttention(in_dim=ch, embed_dim=ch, num_heads=num_heads, p_rope=p_rope, base_rope=base_rope)
 
        self.ff        = FeedForwardBlock(in_dim=ch, hidden_dim=2*ch, dropout=dropout)   
        self.norm_self = nn.RMSNorm(ch)
        self.norm_ff   = nn.RMSNorm(ch)
        self.drop      = nn.Dropout(dropout)

        self._init_weights()
          
    def _init_weights(self) -> None:

        # note a bonus of res-pos-norm is that we can init as identity!
        nn.init.zeros_(self.norm_self.weight)  
        nn.init.zeros_(self.norm_ff.weight)
    
    def forward(self, x: torch.Tensor, pos_idx: torch.Tensor) -> torch.Tensor:
        """
        Assumes batch first.
        
        Shapes:
            x       ... [b, n, ch]       
            pos_idx ... [b, n]
        """

        # Self-attention part
        self_out = x
        self_out = self.self_att(query=self_out, key=self_out, value=self_out, pos_idx=pos_idx)
        self_out = self.norm_self(self_out)  
        self_out = self.drop(self_out) + x      

        # Feed-Forward part
        feed_out = self.ff(self_out)
        feed_out = self.norm_ff(feed_out)
        feed_out = self.drop(feed_out) + self_out                      
        return feed_out

# %% ../../../src/models/clip/unitary_clip.ipynb #1d7888f9-f229-4f75-8599-df4d0e7f202a
class PackingTransformer(ConfigModel):
    """
    The first stage packing/unpacking transformers of the CirDiT model. 
    Applies a RoPE for time dimension only, not on spatial dimension.
    """
    
    def __init__(self,
                 ch: int, 
                 depth: int,
                 num_heads: int, 
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()

        self.blocks = nn.ModuleList([
                        SelfAttnBlock(ch=ch, 
                                      num_heads=num_heads, 
                                      dropout=dropout, 
                                      p_rope=p_rope, 
                                      base_rope=base_rope)
                        for d in range(depth)
                    ])
                
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Shapes:
            x   ... [b, s, t, ch]
        """

        b, s, t, ch = x.shape

        # create pos_idx such that they only depend on the time position
        pos_idx = torch.arange(t, device=x.device, dtype=torch.int32).expand(b, s, -1)
        pos_idx = pos_idx.reshape(b, -1)

        # flatten spatial and time into seq
        x = x.reshape(b, s*t, ch)
        
        for block in self.blocks:
            x = block(x=x, pos_idx=pos_idx)

        # undo flatten
        x = x.reshape(b, s, t, ch)
        
        return x

# %% ../../../src/models/clip/unitary_clip.ipynb #ef895de4-cdd0-4d4e-81d4-1c0ec3bf1341
class CoreTransformer(nn.Module):
    """
    The main transformer of the `CirDiT` model.  
    Applies a RoPE for time dimension.
    """

    def __init__(self,
                 ch: int, 
                 depth: int,
                 num_heads: int, 
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000) -> None:
        super().__init__()
        
        self.blocks = nn.ModuleList([
                        SelfAttnBlock(ch=ch, 
                                      num_heads=num_heads, 
                                      dropout=dropout, 
                                      p_rope=p_rope, 
                                      base_rope=base_rope)
                        for d in range(depth)
                    ])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Shapes:
            x   ... [b,  t, ch]
        """
        
        pos_idx = torch.arange(x.shape[1], device=x.device, dtype=torch.int32) 
        
        for block in self.blocks:
            x = block(x=x, pos_idx=pos_idx)

        return x

# %% ../../../src/models/clip/unitary_clip.ipynb #4b1a77ad-c900-4a0b-8b8e-bf7a2fdddaa4
@dataclass
class CircuitEncoderConfig:
    embedder_config: dict
    
    ch_packing: int 
    ch_core: int
    
    depth_packing: int
    depth_core: int  
    
    num_heads_packing: int
    num_heads_core: int  
    
    dropout: float 
    p_rope: float
    base_rope: float

# %% ../../../src/models/clip/unitary_clip.ipynb #7d534173-c74d-480d-bf36-645ed8ccce9b
class CircuitEncoder(ConfigModel):
    def __init__(self,
                 embedder_config: Optional[dict], 
                 ch_packing: int, 
                 ch_core: int,
                 depth_packing: int,
                 depth_core: int,                 
                 num_heads_packing: int,
                 num_heads_core: int,              
                 dropout: float = 0.0, 
                 p_rope: float = 1.0, 
                 base_rope: float = 10_000,
                 embedder: Optional[nn.Module] = None) -> None:    
        super().__init__()

        if exists(embedder):
            self.embedder = embedder
            embedder_config = self.embedder.get_config(None)
        else:
            assert exists(embedder_config)

        self.params_config = CircuitEncoderConfig(embedder_config=embedder_config,
                                                  ch_packing=ch_packing, 
                                                  ch_core=ch_core,
                                                  depth_packing=depth_packing,
                                                  depth_core=depth_core,                 
                                                  num_heads_packing=num_heads_packing,
                                                  num_heads_core=num_heads_core,              
                                                  dropout=dropout, 
                                                  p_rope=p_rope, 
                                                  base_rope=base_rope)

        if not_exists(embedder):
            self.embedder = ConfigModel.from_config(embedder_config, device=None, silent=True)

        self.packing = PackingTransformer(ch=ch_packing, 
                                          depth=depth_packing, 
                                          num_heads=num_heads_packing, 
                                          dropout=dropout, 
                                          p_rope=p_rope, 
                                          base_rope=base_rope)

        self.core = CoreTransformer(ch=ch_core, 
                                    depth=depth_core, 
                                    num_heads=num_heads_core, 
                                    dropout=dropout, 
                                    p_rope=p_rope, 
                                    base_rope=base_rope)

        self.encoding_ch = ch_core
        
        self.proj_in   = nn.Linear(self.embedder.embedding_dim, ch_packing)
        self.core_proj = nn.Linear(ch_packing, ch_core)

        self.norm_packing = nn.RMSNorm(ch_packing) 
        self.norm_core    = nn.RMSNorm(ch_core) 
        self.norm_final   = nn.RMSNorm(ch_core)
        
        self.qubit_pos_enc = LearnedPositionalEmbedding(dim=ch_packing, max_seq_len=64) #here max number of qubits

        self._init_weights()
   
    def _init_weights(self) -> None:
        nn.init.orthogonal_(self.core_proj.weight)
        nn.init.zeros_(self.core_proj.bias)

    def forward(self, tokens: torch.Tensor, params: torch.Tensor, pool: bool = False) -> torch.Tensor: 
        # Embed the circuits
        x = self.embedder(h=tokens, w=params) 

        # Pre-process circuit and add pos-encoding
        b, s, t, _ = x.shape

        x = self.proj_in(x)
        x = self.qubit_pos_enc(x)

        # Pack spatial into tokens
        x = self.norm_packing(x)  
        x = self.packing(x=x)
        
        # Downsample, reduce spatial, ... [b, t, ch_core]
        x_main = x.mean(dim=1)   
        x_main = self.core_proj(x_main)  
        
        # Core transformer
        x_main = self.norm_core(x_main)
        x_main = self.core(x=x_main)

        if pool: 
            x_main = torch.mean(x_main, dim=1) # [b, ch]  
            
        x_main = self.norm_final(x_main)
        return x_main

# %% ../../../src/models/clip/unitary_clip.ipynb #fbd4627e-3176-475a-91f0-899c18e0f724
@dataclass
class UnitaryCLIPConfig:
    text_encoder_config: dict
    clip_embed_size: int     

# %% ../../../src/models/clip/unitary_clip.ipynb #0b3deb74-e352-438b-9549-a44706596c94
class UnitaryCLIP(ConfigModel):
    
    def __init__(self, 
                 text_encoder_config: Optional[dict],
                 unitary_text_encoder: UnitaryTextEncoder,
                 circuit_encoder: CircuitEncoder,
                 clip_embed_size: int,
                 text_encoder: Optional[nn.Module] = None) -> None:
        super().__init__()        

        if exists(text_encoder):
            self.text_encoder = text_encoder
            text_encoder_config = self.text_encoder.get_config(None)
        else:
            assert exists(text_encoder_config)
        
        self.params_config = UnitaryCLIPConfig(text_encoder_config=text_encoder_config,
                                               clip_embed_size=clip_embed_size)

        if not_exists(text_encoder):
            if "device" in text_encoder_config:
                device = text_encoder_config["device"]
            else:
                device = "cpu"
                
            self.text_encoder = ConfigModel.from_config(text_encoder_config, device=device, silent=True)
            
        self.unitary_text_encoder = unitary_text_encoder
        self.circuit_encoder      = circuit_encoder

        self.unitary_text_proj = nn.Linear(self.unitary_text_encoder.encoding_ch, clip_embed_size)
        self.circuit_proj      = nn.Linear(self.circuit_encoder.encoding_ch     , clip_embed_size)     
        self.temperature       = torch.nn.Parameter(torch.zeros(1))
       
        self._init_weights()
          
    def _init_weights(self) -> None:
        initrange = 0.1   
        self.unitary_text_proj.bias.data.zero_()
        self.unitary_text_proj.weight.data.uniform_(-initrange, initrange)      
        self.circuit_proj.bias.data.zero_()
        self.circuit_proj.weight.data.uniform_(-initrange, initrange)
     
    def forward(self, tokens: torch.Tensor, params: torch.Tensor, y: torch.Tensor, U: torch.Tensor) -> torch.Tensor:

        y_emb = self.text_encoder(y, pool=False)
        
        ut_enc = self.unitary_text_encoder(y_emb=y_emb, U=U, pool=True) 
        ut_enc = self.unitary_text_proj(ut_enc)   # out [b, embed_size]
        ut_enc = F.normalize(ut_enc, dim=-1)  
        
        #------------  
        
        qc_enc = self.circuit_encoder(tokens=tokens, params=params, pool=True) 
        qc_enc = self.circuit_proj(qc_enc)       # out [b, embed_size]      
        qc_enc = F.normalize(qc_enc, dim=-1)   
        
        #------------     
        
        scores = torch.matmul(ut_enc, qc_enc.T) * torch.exp(self.temperature) #[b, b]
         
        #scores is: I=unitary_text   T=circuit
        #--------------------------------
        #| I1*T1   I1*T2   I1*T3   ...
        #| I2*T1
        #| I3*T1
        # ...
        #--------------------------------
            
        return scores
