Method test_constant

python/tests/test_init.py:11–20

(self)

class TestInit(mlx_tests.MLXTestCase):
    def test_constant(self):
        value = 5.0

        for dtype in [mx.float32, mx.float16]:
            initializer = init.constant(value, dtype)
            for shape in [(3,), (3, 3), (3, 3, 3)]:
                result = initializer(mx.array(mx.zeros(shape)))
                with self.subTest(shape=shape):
                    self.assertEqual(result.shape, shape)
                    self.assertEqual(result.dtype, dtype)

    def test_normal(self):
        mean = 0.0