Class CrossAttnEncoder

src/models/text_decoder_v2.py:285–358

A Transformer Encoder with interleaved self-attention and cross-attention. This architecture processes a primary sequence `x` (e.g., text) by alternating between self-attention blocks (to process `x` itself) and cross-attention blocks (to infuse information from a context sequence `u`,

class CrossAttnEncoder(nn.Module):
    """A Transformer Encoder with interleaved self-attention and cross-attention.

    This architecture processes a primary sequence `x` (e.g., text) by alternating
    between self-attention blocks (to process `x` itself) and cross-attention
    blocks (to infuse information from a context sequence `u`, e.g., an image).

    Attributes:
        depth: The total number of layers. Each "layer" consists of one
               self-attention block and one cross-attention block.
        remat_policy: Gradient checkpointing policy ('none', 'minimal', etc.).
    """
    depth: int
    mlp_dim: Optional[int] = None
    num_heads: int = 12
    dropout: float = 0.0
    drop_path: float = 0.0
    remat_policy: str = "none"
    casual_mask: bool = False
    use_flash_attn: bool = False
    dtype: Optional[Dtype] = jnp.float32
    param_dtype: Dtype = jnp.float32
    mesh: Optional[Any] = None

    @nn.compact
    def __call__(self, x: Array, u: Array, deterministic: bool = True) -> Tuple[Array, dict]:
        out = {}
        # Linearly increasing drop path rate for stochastic depth
        dpr = [rate.item() for rate in np.linspace(0, self.drop_path, self.depth)]
        
        # Configure gradient checkpointing (remat) if specified
        if self.remat_policy not in (None, "none"):
            policy = (jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
                      if self.remat_policy == "minimal" else None)
            logging.info(f"Applying activation checkpointing: {self.remat_policy}")
            SelfAttnBlock = nn.remat(
                Encoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(1,))
            CrossAttnBlock = nn.remat(
                CrossAttnEncoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(2,))
        else:
            SelfAttnBlock = Encoder1DBlock
            CrossAttnBlock = CrossAttnEncoder1DBlock
            
        for i in range(self.depth):
            # 1. Self-attention on text embeddings `x`
            x, out[f"self_attn_block_{i:02d}"] = SelfAttnBlock(
                name=f"self_attn_block_{i}",
                mlp_dim=self.mlp_dim,
                depth=self.depth,
                num_heads=self.num_heads,
                dropout=self.dropout,
                drop_path=dpr[i],
                casual_mask=self.casual_mask,
                use_flash_attn=self.use_flash_attn,
                dtype=self.dtype,
                param_dtype=self.param_dtype,
                mesh=self.mesh
            )(x, deterministic)