Class CrossAttnEncoder1DBlock

src/models/text_decoder_v2.py:208–282

A Transformer block with cross-attention followed by an MLP. This block takes two inputs: `x` (typically text embeddings) which acts as the query, and `u` (typically image embeddings) which acts as the key and value for the cross-attention layer. Attributes: mlp_dim: The hi

class CrossAttnEncoder1DBlock(nn.Module):
    """A Transformer block with cross-attention followed by an MLP.

    This block takes two inputs: `x` (typically text embeddings) which acts
    as the query, and `u` (typically image embeddings) which acts as the
    key and value for the cross-attention layer.

    Attributes:
        mlp_dim: The hidden dimension of the MLP block.
        num_heads: The number of attention heads.
        dropout: Dropout rate for attention and MLP outputs.
        drop_path: Stochastic depth drop rate.
        depth: The total number of layers in the parent encoder, used for weight init scaling.
        use_flash_attn: Whether to use TPU Flash Attention.
        dtype: The computation data type.
        param_dtype: The parameter data type.
        mesh: The JAX device mesh for model parallelism.
    """
    mlp_dim: Optional[int] = None
    num_heads: int = 12
    dropout: float = 0.0
    drop_path: float = 0.0
    depth: int = 12
    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]:
        width = x.shape[-1]
        # Scaled weight initialization based on model depth
        init_std = {
            'proj': (width ** -0.5) * ((2 * self.depth) ** -0.5),
            'attn': width ** -0.5,
            'fc': (2 * width) ** -0.5
        }
        out = {}

        # Cross-Attention part
        y = nn.LayerNorm(dtype=self.dtype, param_dtype=self.param_dtype, name="ln_x")(x)
        v = nn.LayerNorm(dtype=self.dtype, param_dtype=self.param_dtype, name="ln_u")(u)
        
        y = out["cross_attn"] = MultiHeadDotProductAttention(
            num_heads=self.num_heads,
            attn_kernel_init=nn.initializers.normal(stddev=init_std['attn']),
            proj_kernel_init=nn.initializers.normal(stddev=init_std['proj']),
            bias_init=nn.initializers.zeros,
            dtype=self.dtype,
            param_dtype=self.param_dtype,
            deterministic=deterministic,
            use_flash_attn=self.use_flash_attn,
            mesh=self.mesh,
        )(inputs_q=y, inputs_kv=v)
        
        y = nn.Dropout(rate=self.dropout)(y, deterministic)
        y = DropPath(dropout_prob=self.drop_path)(y, deterministic)
        x = x + y