(jt, padding_val)
| 297 | ###################################################################### |
| 298 | # Generate padded forms of query, key, value for comparison |
| 299 | def jagged_to_padded(jt, padding_val): |
| 300 | # TODO: do jagged -> padded directly when this is supported |
| 301 | return torch.nested.to_padded_tensor( |
| 302 | torch.nested.nested_tensor(list(jt.unbind())), |
| 303 | padding_val) |
| 304 | |
| 305 | padded_query, padded_key, padded_value = ( |
| 306 | jagged_to_padded(t, 0.0) for t in (query, key, value) |