Class CrossAttnEncoder1DBlock

src/models/text_decoder.py:249–331

Single transformer encoder block (MHSA + MLP).

class CrossAttnEncoder1DBlock(nn.Module):
    """Single transformer encoder block (MHSA + MLP)."""
    mlp_dim: Optional[int] = None  # Defaults to 4x input dim
    num_heads: int = 12
    dropout: float = 0.0
    drop_path: float = 0.0
    depth: int = 12
    casual_mask: bool = False
    use_flash_attn: bool = False
    dtype: Optional[Dtype] = jnp.float32
    param_dtype: Dtype = jnp.float32
    mesh: Any = None

    @nn.compact
    def __call__(self, x, u, attn_mask=None, deterministic=True):
        width = x.shape[-1]
        init_std = {
            'proj': (width ** -0.5) * ((2 * self.depth) ** -0.5),
            'attn': width ** -0.5,
            'fc': (2 * width) ** -0.5
        }
        out = {}
        x = x.astype(self.dtype)
        x = nn.with_logical_constraint(x, ("activation_batch", "activation_length", "activation_embed"))
        y = nn.LayerNorm(
            dtype=self.dtype,
            param_dtype=self.param_dtype,
            scale_init=nn.with_logical_partitioning(nn.initializers.ones_init(), ("norm",)),
            bias_init=nn.with_logical_partitioning(nn.initializers.zeros_init(), (None,)),
           )(x)
        y = nn.with_logical_constraint(y, ("activation_batch", "activation_length", "activation_embed"))

        u = u.astype(self.dtype)
        u = nn.with_logical_constraint(u, ("activation_batch", "activation_length", "activation_embed"))
        v = nn.LayerNorm(
            dtype=self.dtype,
            param_dtype=self.param_dtype,
            scale_init=nn.with_logical_partitioning(nn.initializers.ones_init(), ("norm",)),
            bias_init=nn.with_logical_partitioning(nn.initializers.zeros_init(), (None,)),
           )(u)
        v = nn.with_logical_constraint(v, ("activation_batch", "activation_length", "activation_embed"))
        
        y = out["sa"] = 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,
            scan_attn=self.scan_attn,
            scan_attn_chunck=self.mlp_chunck,
            mesh=self.mesh,
        )(y, v, mask=attn_mask if self.casual_mask else None)
        y = nn.Dropout(rate=self.dropout)(y, deterministic)
        # y = DropPath(dropout_prob=self.drop_path)(y, deterministic)
        x = out["+sa"] = x + y