(self)
| 1972 | self.assertEqual(c_out.shape, (44, 12)) |
| 1973 | |
| 1974 | def test_quantized_embedding(self): |
| 1975 | emb = nn.Embedding(32, 256) |
| 1976 | qemb = nn.QuantizedEmbedding.from_embedding(emb, bits=8) |
| 1977 | x = mx.array([2, 6, 9, 3, 0, 3]) |
| 1978 | y = emb(x) |
| 1979 | yq = qemb(x) |
| 1980 | self.assertLess((y - yq).abs().max(), qemb.scales.max()) |
| 1981 | |
| 1982 | x = mx.random.uniform(shape=(2, 256)) |
| 1983 | y = emb.as_linear(x) |
| 1984 | yq = qemb.as_linear(x) |
| 1985 | |
| 1986 | def cosine(a, b): |
| 1987 | ab = (a * b).sum(-1) |
| 1988 | aa = mx.linalg.norm(a, axis=-1) |
| 1989 | bb = mx.linalg.norm(b, axis=-1) |
| 1990 | return ab / aa / bb |
| 1991 | |
| 1992 | self.assertGreater(cosine(y, yq).min(), 0.99) |
| 1993 | |
| 1994 | def test_causal_mask(self): |
| 1995 | mask = nn.MultiHeadAttention.create_additive_causal_mask(4, mx.float16) |
nothing calls this directly
no test coverage detected