Class Encoder1DBlock

src/models/vit.py:245–334

Single transformer encoder block (MHSA + MLP).

class Encoder1DBlock(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
  init_values: float = None
  use_flash_attn: bool = False
  dtype: Optional[Dtype] = jnp.float32
  param_dtype: Dtype = jnp.float32
  mesh: Any = None
  use_dense_general: bool = False
  scan_mlp: bool = False
  scan_attn: bool = False
  mlp_chunck: int = 128

  @nn.compact
  def __call__(self, x, deterministic=True):
    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"))

    ## hack init func
    y = out["sa"] = common.MultiHeadDotProductAttention(
        num_heads=self.num_heads,
        qkv_kernel_init=nn.initializers.normal(stddev=0.02),
        out_kernel_init=nn.initializers.normal(stddev=0.02),
        bias_init=nn.initializers.zeros,
        deterministic=deterministic,
        use_flash_attn=self.use_flash_attn,
        scan_attn=self.scan_attn,
        scan_attn_chunck=self.mlp_chunck,
        dtype=self.dtype,
        param_dtype=self.param_dtype,
        mesh=self.mesh,
        use_dense_general=self.use_dense_general
    )(y, y)
    y = nn.with_logical_constraint(y, ("activation_batch", "activation_length", "activation_embed"))

    y = nn.Dropout(rate=self.dropout)(y, deterministic)
    if self.init_values is not None:
        n, l, d = y.shape
        y = LayerScale(d, init_values=self.init_values, name='ls1')(y)
    y = DropPath(dropout_prob=self.drop_path)(y, deterministic)
    x = out["+sa"] = x + y
    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",)),