diff --git a/mlx/random.cpp b/mlx/random.cpp index 63e39cdcc..5e0682d32 100644 --- a/mlx/random.cpp +++ b/mlx/random.cpp @@ -153,14 +153,23 @@ array uniform( array normal( const std::vector& shape, Dtype dtype, + const float loc /* = 0.0 */, + const float scale /* = 1.0 */, const std::optional& key /*= nullopt */, StreamOrDevice s /* = {} */) { auto stream = to_stream(s); auto low = array(std::nextafter(-1.0f, 0.0f), dtype); auto high = array(1.0f, dtype); auto samples = uniform(low, high, shape, dtype, key, stream); - return multiply( - array(std::sqrt(2.0), dtype), erfinv(samples, stream), stream); + samples = + multiply(array(std::sqrt(2.0), dtype), erfinv(samples, stream), stream); + if (scale != 1.0) { + samples = multiply(array(scale, dtype), samples, stream); + } + if (loc != 0.0) { + samples = add(array(loc, dtype), samples, stream); + } + return samples; } array randint( diff --git a/mlx/random.h b/mlx/random.h index ab75eb488..1397b32d7 100644 --- a/mlx/random.h +++ b/mlx/random.h @@ -95,13 +95,30 @@ inline array uniform( array normal( const std::vector& shape, Dtype dtype, + const float loc, + const float scale, const std::optional& key = std::nullopt, StreamOrDevice s = {}); inline array normal( const std::vector& shape, + const float loc, + const float scale, const std::optional& key = std::nullopt, StreamOrDevice s = {}) { - return normal(shape, float32, key, s); + return normal(shape, float32, loc, scale, key, s); +} +inline array normal( + const std::vector& shape, + const Dtype dtype, + const std::optional& key = std::nullopt, + StreamOrDevice s = {}) { + return normal(shape, dtype, 0.0, 1.0, key, s); +} +inline array normal( + const std::vector& shape, + const std::optional& key = std::nullopt, + StreamOrDevice s = {}) { + return normal(shape, float32, 0.0, 1.0, key, s); } /** Generate integer samples uniformly at random */ diff --git a/python/src/random.cpp b/python/src/random.cpp index 6e9f38d97..e9140e7d9 100644 --- a/python/src/random.cpp +++ b/python/src/random.cpp @@ -99,13 +99,17 @@ void init_random(py::module_& parent_module) { "normal", [](const std::vector& shape, std::optional type, + float loc, + float scale, const std::optional& key, StreamOrDevice s) { - return normal(shape, type.value_or(float32), key, s); + return normal(shape, type.value_or(float32), loc, scale, key, s); }, "shape"_a = std::vector{}, "dtype"_a = std::optional{float32}, + "loc"_a = 0.0, + "scale"_a = 1.0, "key"_a = none, "stream"_a = none, R"pbdoc( @@ -114,6 +118,8 @@ void init_random(py::module_& parent_module) { Args: shape (list(int), optional): Shape of the output. Default is ``()``. dtype (Dtype, optional): Type of the output. Default is ``float32``. + loc (float, optional): Mean of the distribution. Default is ``0.0``. + scale (float, optional): Standard deviation of the distribution. Default is ``1.0``. key (array, optional): A PRNG key. Default: None. Returns: diff --git a/python/tests/test_random.py b/python/tests/test_random.py index 0a06c3496..892db37df 100644 --- a/python/tests/test_random.py +++ b/python/tests/test_random.py @@ -80,6 +80,20 @@ class TestRandom(mlx_tests.MLXTestCase): a = mx.random.normal(dtype=t) self.assertEqual(a.dtype, t) + # Generate with a given mean and standard deviation + loc = 1.0 + scale = 2.0 + + a = mx.random.normal(shape=(3, 2), loc=loc, scale=scale, key=key) + b = scale * mx.random.normal(shape=(3, 2), key=key) + loc + self.assertTrue(mx.allclose(a, b)) + + a = mx.random.normal( + shape=(3, 2), loc=loc, scale=scale, dtype=mx.float16, key=key + ) + b = scale * mx.random.normal(shape=(3, 2), dtype=mx.float16, key=key) + loc + self.assertTrue(mx.allclose(a, b)) + self.assertEqual(mx.random.normal().dtype, mx.random.normal(dtype=None).dtype) def test_randint(self):