Pooling implementation is weird
Hi,
I'm interested in your implementation of checking for left padding. Isn't this incorrect in the case all the inputs are of equal length for right side padding?
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
This is much easier to check using the tokenizer:
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6-Embedding")
print(tok.padding_side) # "right"
Equal-length right-padded inputs do not break this shortcut when the batch has no padding. The name left_padding is misleading: for binary masks, the test asks whether every row's final position is valid.
In the pinned function, that condition selects [:, -1]. For unpadded rows, sum(mask)-1 names the same final position. If equal-length rows are padded to a longer fixed width on the right, the condition is false and the length-based branch skips the pads.
This runnable illustration uses token positions, not embeddings:
for mask in ([[1, 1, 1], [1, 1, 1]],
[[1, 1, 0], [1, 1, 0]]):
final_valid = all(row[-1] == 1 for row in mask)
positions = [len(row) - 1 if final_valid else sum(row) - 1
for row in mask]
print(final_valid, positions)
# True [2, 2]
# False [1, 1]
The configured padding side and the current batch's mask are different facts. Keep one padding side per batch and at least one valid token per row; mixed sides or masks with gaps are outside this shortcut.
I checked synthetic index arithmetic using a NumPy adaptation of the function. PyTorch, tokenizer, model inference and retrieval quality were not tested.
CyberNative AI LLC 路 AI-assisted technical note; corrections: hello@cybernative.ai.