Gradients for Tied Token Embeddings

Problem

A tiny tied language model looks up hidden = embedding[token_ids], then computes logits = hidden @ embedding.T. Given dlogits, return the total derivative of the scalar loss with respect to embedding.

The same matrix participates in both the input lookup and output projection. Include both contributions, and accumulate correctly when token_ids contains repeats. No softmax or transformer is part of this problem.

def tied_embedding_gradient(embedding, token_ids, dlogits):

Constraints

• embedding has shape V × D, token_ids length T, dlogits shape T × V.
• 1 ≤ V,D,T ≤ 8; every token ID is in [0,V).
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      2
    ]
  ],
  [
    0
  ],
  [
    [
      1
    ]
  ]
]
Expected result: [
  [
    4.000000000026205
  ]
]

Arguments: [
  [
    [
      1
    ],
    [
      2
    ]
  ],
  [
    0,
    0
  ],
  [
    [
      1,
      0
    ],
    [
      0,
      1
    ]
  ]
]
Expected result: [
  [
    4.000000000004
  ],
  [
    1.0000000000065512
  ]
]