zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 01:50: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: