zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-18 15:28:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add loc and scale to random.normal (#638)

Browse Source 
* Add loc and scale to random.normal

* Add tests for loc and scale for random.normal

* Run pre-commit hooks

* Fix code review
This commit is contained in:
Noah Farr
2024-02-07 20:49:59 +01:00
committed by GitHub
parent ef73393a19
commit 5fd11c347d
4 changed files with 50 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/src/random.cpp
View File

@@ -99,13 +99,17 @@ void init_random(py::module_& parent_module) {
"normal",
[](const std::vector<int>& shape,
std::optional<Dtype> type,
float loc,
float scale,
const std::optional<array>& 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<int>{},
"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:

14
python/tests/test_random.py
View File

@@ -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):