expose prob as positional argument for bernoulli ops (#1949)

mindspore-lab · Feb 8, 2025 · 9ea7a99 · 9ea7a99
1 parent 6c357b8
commit 9ea7a99
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mindnlp/core/ops/random.py b/mindnlp/core/ops/random.py
@@ -11,11 +11,11 @@
 
 # bernoulli
 has_bernoulli = hasattr(mindspore.mint, 'bernoulli')
-def bernoulli(input, *, generator=None):
+def bernoulli(input, *, generator=None, p=0.5):
     if use_pyboost() and has_bernoulli:
         return mindspore.mint.bernoulli(input, generator=generator)
     random_numbers = rand(*input.shape, dtype=mindspore.float32)
-    samples = random_numbers < 0.5
+    samples = random_numbers < p
     samples = samples.int()
     return samples