Method test_multivariate_normal

python/tests/test_random.py:112–200

(self)

                self.assertTrue(mx.all(a < mx.inf))

    def test_multivariate_normal(self):
        key = mx.random.key(0)
        mean = mx.array([0, 0])
        cov = mx.array([[1, 0], [0, 1]])

        a = mx.random.multivariate_normal(mean, cov, key=key, stream=mx.cpu)
        self.assertEqual(a.shape, (2,))

        ## Check dtypes
        for t in [mx.float32]:
            a = mx.random.multivariate_normal(
                mean, cov, dtype=t, key=key, stream=mx.cpu
            )
            self.assertEqual(a.dtype, t)
        for t in [
            mx.int8,
            mx.int32,
            mx.int64,
            mx.uint8,
            mx.uint32,
            mx.uint64,
            mx.float16,
            mx.bfloat16,
        ]:
            with self.assertRaises(ValueError):
                mx.random.multivariate_normal(
                    mean, cov, dtype=t, key=key, stream=mx.cpu
                )

        ## Check incompatible shapes
        with self.assertRaises(ValueError):
            mean = mx.zeros((2, 2))
            cov = mx.zeros((2, 2))
            mx.random.multivariate_normal(mean, cov, shape=(3,), key=key, stream=mx.cpu)

        with self.assertRaises(ValueError):
            mean = mx.zeros((2))
            cov = mx.zeros((2, 2, 2))
            mx.random.multivariate_normal(mean, cov, shape=(3,), key=key, stream=mx.cpu)

        with self.assertRaises(ValueError):
            mean = mx.zeros((3,))
            cov = mx.zeros((2, 2))
            mx.random.multivariate_normal(mean, cov, key=key, stream=mx.cpu)

        with self.assertRaises(ValueError):
            mean = mx.zeros((2,))
            cov = mx.zeros((2, 3))
            mx.random.multivariate_normal(mean, cov, key=key, stream=mx.cpu)

        ## Different shape of mean and cov
        mean = mx.array([[0, 7], [1, 2], [3, 4]])
        cov = mx.array([[1, 0.5], [0.5, 1]])
        a = mx.random.multivariate_normal(mean, cov, shape=(4, 3), stream=mx.cpu)
        self.assertEqual(a.shape, (4, 3, 2))

        ## Check correcteness of the mean and covariance
        n_test = int(1e5)