How many sublayers per layer?

Background

The two stacks differ by one sublayer:

• Encoder layer: self-attention → FFN → 2 sublayers.
• Decoder layer: masked self-attention → cross-attention → FFN → 3 sublayers.

Each sublayer is wrapped in its own Add & Norm.

Problem statement

Implement sublayers_per_layer(stack) returning the number of sublayers.

Input

• stack — "encoder" or "decoder".

Output

Returns an int: 2 for "encoder", 3 for "decoder".

Examples

Example 1

Input:  stack = "encoder"
Output: 2

Example 2

Input:  stack = "decoder"
Output: 3

Constraints

• "encoder" → 2; "decoder" → 3.
• Return a plain int.

Notes

• The extra decoder sublayer is cross-attention — the only place the decoder reads the encoder's memory.