Project and Split Packed QKV

Problem

Project token rows through a packed QKV matrix. x is T × I and weight is I × (3D); its columns contain all Q features, then all K, then all V. Return {"q": Q, "k": K, "v": V}, each of shape heads × T × (D/heads). Within each projection, consecutive features belong to consecutive heads.

def packed_qkv(x, weight, heads):

Constraints

• 1 ≤ T,I,D ≤ 12; heads is a positive divisor of D.
• Inputs are finite; no bias is applied.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      1
    ]
  ],
  [
    [
      1,
      2,
      3
    ]
  ],
  1
]
Expected result: {
  "q": [
    [
      [
        1
      ]
    ]
  ],
  "k": [
    [
      [
        2
      ]
    ]
  ],
  "v": [
    [
      [
        3
      ]
    ]
  ]
}

Arguments: [
  [
    [
      1,
      2
    ]
  ],
  [
    [
      1,
      0,
      0,
      0,
      0,
      0
    ],
    [
      0,
      1,
      0,
      0,
      0,
      0
    ]
  ],
  2
]
Expected result: {
  "q": [
    [
      [
        1
      ]
    ],
    [
      [
        2
      ]
    ]
  ],
  "k": [
    [
      [
        0
      ]
    ],
    [
      [
        0
      ]
    ]
  ],
  "v": [
    [
      [
        0
      ]
    ],
    [
      [
        0
      ]
    ]
  ]
}