Spaces:
Running
Running
| """Numerical stability of softmax / cosine.""" | |
| from __future__ import annotations | |
| import numpy as np | |
| from lightweight_embeddings.core import math_ops | |
| def test_softmax_sums_to_one(): | |
| a = np.array([[1.0, 2.0, 3.0], [10.0, 0.0, -10.0]]) | |
| out = math_ops.softmax(a) | |
| assert np.allclose(out.sum(axis=-1), 1.0) | |
| def test_softmax_handles_large_values(): | |
| a = np.array([1e6, 1e6 + 1.0]) | |
| out = math_ops.softmax(a) | |
| assert np.isfinite(out).all() | |
| def test_cosine_similarity_identity(): | |
| a = np.eye(3, dtype=np.float32) | |
| sim = math_ops.cosine_similarity(a, a) | |
| assert np.allclose(sim, np.eye(3), atol=1e-6) | |
| def test_cosine_similarity_zero_vectors_safe(): | |
| a = np.zeros((1, 4), dtype=np.float32) | |
| b = np.zeros((1, 4), dtype=np.float32) | |
| sim = math_ops.cosine_similarity(a, b) | |
| # Should not be NaN — clamp to zero norms keeps result finite. | |
| assert np.isfinite(sim).all() | |
| def test_normalize_unit_norm(): | |
| a = np.random.RandomState(0).randn(5, 7).astype(np.float32) | |
| out = math_ops.normalize(a, axis=1) | |
| assert np.allclose(np.linalg.norm(out, axis=1), 1.0, atol=1e-6) | |