Masked Attention with PyTorch | Dataford Interview Questions

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 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 question is Masked Attention with PyTorch. Start with the requirements in the Question tab.

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

Problem

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

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 question is Masked Attention with PyTorch. Start with the requirements in the Question tab.

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