Masked Attention with PyTorch

HardPython00:00

Practice interviewer

In session

5 left

00:00

Your question is Masked Attention with PyTorch. Start with the requirements on the right.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.

You need to log in / sign up to run or submit.

Problem

Write a script using PyTorch tensors to perform custom masked attention operations for a sequence-to-sequence model.

For dependency-free grading, implement the same operation over tensor-shaped nested Python lists. The inputs use shapes q, k, v = [B][H][Q or K][D] and mask = [B][Q][K]; True permits attention and False blocks it. Return an output with shape [B][H][Q][D], using scaled dot-product attention, softmax over keys, and an all-zero output for a query whose keys are all masked.

Example: q=[[[[0.0]]]], k=[[[[1.0],[2.0]]]], v=[[[[5.0],[9.0]]]], mask=[[[true,true]]] returns [[[[7.0]]]]. If the mask is [[[false,false]]], the result is [[[[0.0]]]].

Constraints

1 <= B, H, Q, K <= 256
1 <= D <= 128
q and k contain numeric values
v has shape [B][H][K][D]
mask has shape [B][Q][K]
Every mask entry is boolean
At least one input query may have all keys masked

Function Signature

def masked_attention(q, k, v, mask):

Your solutionPython 3

You need to log in / sign up to run or submit.

Run your code to see test output