Looped Transformer: Shared-Weight Unrolling

Easy · Solved in PyTorch · Machine learning coding practice

Problem

Implement the depth loop of a looped Transformer. Reuse the supplied block callable exactly num_loops times, feeding each result into the next pass. The block already contains the model operations and shared weights; do not create new blocks or add extra residual connections.

hidden is a floating-point tensor/array of shape (B, S, D) with positive dimensions. block(state) returns the same shape without modifying its input. The integer num_loops is between 0 and 16.

Topics: transformer, looped-transformer, pytorch, numpy

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Looped Transformer: Shared-Weight Unrolling in the browser and run it against the tests.

Related problems