THE MULTI-HEAD DIMENSION CHECK

Difficulty: Hard Tags: Transformers, Multi-Head Attention, Tensor Shapes

Story

A transformer implementation splits a hidden vector into attention heads. A configuration loader accepts a hidden size that is not divisible by the number of heads, leading to silent reshape errors. Validate the configuration and compute head dimension and projection parameter count.

The Math

For hidden size D and head count H, the configuration is valid iff D mod H == 0. Then head dimension is d_h = D/H. For bias-free query, key, value, and output projections, each mapping D -> D, total projection parameters are:

N_params = 4 * D^2

Head count changes partitioning, not this parameter count.

Input Format

D (int), H (int).

Output Format

None if D is not divisible by H. Otherwise (head_dim, total_params).

Constraints

• 1 <= D, H <= 10^9. Use wide integer arithmetic.
• Time limit: 1.0 second.

Example

Input

D = 768, H = 12

Output

(64, 2359296)

Explanation: The head dimension is 768/12 = 64. Four dense 768x768 projections contain 4*768^2 = 2359296 weights.