add trellis quant mode

2025-11-05 11:28:12 +08:00 · 2025-03-18 18:52:22 -07:00
parent e9e268336b
commit d7acf59fd0
16 changed files with 852 additions and 108 deletions
--- a/python/tests/test_quantized.py
+++ b/python/tests/test_quantized.py
@@ -10,6 +10,9 @@ import mlx_tests
 class TestQuantized(mlx_tests.MLXTestCase):
    def test_quantize_dequantize(self):
        w = mx.random.normal(shape=(128, 512))
+        w_q, scales, biases = mx.quantize(w, bits=2, mode="trellis")
+        print(w_q, scales, biases)
+
        for gs in [32, 64, 128]:
            for b in [2, 3, 6, 4, 8]:
                with self.subTest(gs=gs, b=b):