mlx-lm/tests/test_losses.py at 85b94d70c10880fdc6258a851ad36c14292a9b0b · ml-explore/mlx-lm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
# Copyright © 2025 Apple Inc.

import unittest

import mlx.core as mx

from mlx_lm.tuner.losses import can_run_metal, js_div_loss, kl_div_loss


class TestLosses(unittest.TestCase):

    def test_kl_div_loss(self):
        self.assertTrue(can_run_metal())
        mx.random.seed(0)

        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)

        with mx.stream(mx.cpu):
            expected = kl_div_loss(logits_q, logits_p)
        kl = kl_div_loss(logits_q, logits_p)

        self.assertTrue(mx.allclose(kl, expected, rtol=1e-4))

    def test_js_div_loss(self):
        self.assertTrue(can_run_metal())
        mx.random.seed(0)

        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)

        with mx.stream(mx.cpu):
            expected = js_div_loss(logits_q, logits_p)
        js = js_div_loss(logits_q, logits_p)

        self.assertTrue(mx.allclose(js, expected))

    def test_kl_div_loss_vjp(self):
        self.assertTrue(can_run_metal())
        mx.random.seed(0)

        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)

        with mx.stream(mx.cpu):
            expected = mx.vjp(kl_div_loss, [logits_q, logits_p], [cotan])[1][0]
        vjp_q = mx.vjp(kl_div_loss, [logits_q, logits_p], [cotan])[1][0]

        self.assertTrue(mx.allclose(vjp_q, expected))

    def test_js_div_loss_vjp(self):
        self.assertTrue(can_run_metal())
        mx.random.seed(0)

        logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
        logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
        cotan = mx.random.uniform(shape=(4, 8), dtype=mx.float32)

        with mx.stream(mx.cpu):
            expected = mx.vjp(js_div_loss, [logits_q, logits_p], [cotan])[1][0]
        vjp_q = mx.vjp(js_div_loss, [logits_q, logits_p], [cotan])[1][0]

        self.assertTrue(mx.allclose(vjp_q, expected))


if __name__ == "__main__":
    unittest.main()