Repository navigation
fix: normalization underflow in fp16 - #412
Conversation
[Medium impact] The PR appears safe to merge, with one non-blocking memory-use improvement. Reviews (2) · Last reviewed commit: "fix: normalization underflow in fp16" · Reviewed by Greptile |
Codecov Report✅ All modified and coverable lines are covered by tests.
🚀 New features to boost your workflow:
|
| if weights is None: | ||
| weights = cast(np.ndarray, np.linalg.norm(embeddings, axis=1) + 1e-32) | ||
| float_dtype = embeddings.dtype if np.issubdtype(embeddings.dtype, np.floating) else np.dtype(np.float64) | ||
| compute_embeddings = embeddings.astype(np.result_type(embeddings.dtype, np.float32)) |
There was a problem hiding this comment.
Clustering holds an extra matrix
For float32 embeddings without weights, astype copies the entire matrix even though its dtype is unchanged. compute_embeddings then keeps that extra matrix alive throughout KMeans.fit, adding avoidable memory use for large vocabularies. Use copy=False for the initial conversion and release the working array after normalization.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
closes #411