Class Encoder

src/models/text_transformer.py:509–586

Transformer Model Encoder for sequence to sequence translation.

class Encoder(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"
    casual_mask: bool = False
    scan_mlp: bool = False
    scan_attn: bool = False
    mlp_chunck: int = 128
    use_flash_attn: bool = False
    dtype: Optional[Dtype] = jnp.float32
    param_dtype: Dtype = jnp.float32
    mesh: Any = None
    fusion_style: str = "cross_attn"
    li: int = 0
    lt: int = 0
    use_dense_general: bool = False

    @nn.compact
    def __call__(self, x, deterministic=True):
        out = {}
        dpr = [
            float(x) for x in np.linspace(
                0,
                self.drop_path,
                self.depth)]  # drop path decay
        # Input Encoder
        if self.remat_policy != "none":
            if self.remat_policy == "minimal":
                policy = jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
            elif self.remat_policy == "minimal_offloaded":
                policy = jax.checkpoint_policies.offload_dot_with_no_batch_dims(offload_src="device",
                                                                                offload_dst="pinned_host")
            elif self.remat_policy == "minimal_flash":
                policy = jax.checkpoint_policies.save_from_both_policies(
                    jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims,
                    jax.checkpoint_policies.save_only_these_names(
                        "context",
                    ),
                )
            else:
                assert self.remat_policy == "full", "Remat policy needs to be on list of remat policies"
                policy = None

            BlockLayer = remat(  # pylint: disable=invalid-name
                Encoder1DBlock, prevent_cse=True, policy=policy, static_argnums=(1,)
            )  # "deterministic" is a static argu
        else:
            BlockLayer = Encoder1DBlock

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