Weight tying (share input/output embedding matrix)

The problem, from first principles

[02-modern-transformer-architecture/09-untied-embeddings] established the mechanism and the direct contrast in isolation: an input embedding table and an output projection matrix, both shaped (vocab_size, d_model), can either be TWO separate, independently-learned parameters ("untied"), or literally the SAME matrix, consulted in both directions ("tied"). This question wires that choice directly into a full language model's assembly, as a single tied flag controlling which behavior [01-full-forward-pass], immediately following, actually uses.

From theory to code

Implement compute_output_logits(hidden_states, embedding_table, tied, output_weight), dispatching to embedding_table when tied=True (ignoring output_weight entirely) or to output_weight when tied=False, and count_output_head_parameters(vocab_size, d_model, tied), the number of ADDITIONAL parameters the choice introduces.

Constraints

• tied=True: use embedding_table for the projection; output_weight is not needed and may be None.
• tied=False: use output_weight for the projection.
• count_output_head_parameters returns 0 when tied (no new matrix at all) and vocab_size * d_model when untied (one full new matrix).

Hints

matrix = embedding_table if tied else output_weight
return hidden_states @ matrix.T