Class CrossAttnEncoder

src/models/text_decoder.py:335–411

Transformer Model Encoder for sequence to sequence translation.

from flax.linen.partitioning import remat

class CrossAttnEncoder(nn.Module):
    """Transformer Model Encoder for sequence to sequence translation."""
    depth: int
    mlp_dim: Optional[int] = None  # Defaults to 4x input dim
    num_heads: int = 12
    dropout: float = 0.0
    # drop_path: float = 0.0
    remat_policy: str = "none"
    drop_path: float = 0.0
    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, deterministic=True):
        out = {}
        dpr = [
            float(x) for x in np.linspace(
                0,
                self.drop_path,
                self.depth)]  # drop path decay
        # Input Encoder
        CrossAttnBlockLayer = CrossAttnEncoder1DBlock
        if self.remat_policy not in (None, "none"):
            logging.info(f"remat policy: {self.remat_policy}")
            if self.remat_policy == "minimal":
                policy = jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
            else:
                policy = None
            logging.info(f"activation checkpointing {self.remat_policy}")
            CrossAttnBlockLayer = remat(  # pylint: disable=invalid-name
                CrossAttnEncoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(3,)
            )  # "deterministic" is a static argument in CrossAttnEncoder1DBlock

        BlockLayer = Encoder1DBlock
        if self.remat_policy not in (None, "none"):
            logging.info(f"remat policy: {self.remat_policy}")
            if self.remat_policy == "minimal":
                policy = jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
            else:
                policy = None
            logging.info(f"activation checkpointing {self.remat_policy}")
            BlockLayer = remat(  # pylint: disable=invalid-name
                Encoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(1,)
            )  # "deterministic" is a static argument in Encoder1DBlock

        for lyr in range(self.depth):
            x, out[f"block{lyr:02d}"] = BlockLayer(
                name=f"encoderblock_{lyr}",
                mlp_dim=self.mlp_dim,
                depth=self.depth,
                num_heads=self.num_heads,
                dropout=self.dropout,
                drop_path=dpr[lyr],
                casual_mask=self.casual_mask,
                use_flash_attn=self.use_flash_attn,