From 26146e5ab1e0f8d4c0091436125ab6ff127bc28e Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Thu, 23 Dec 2021 13:42:21 +0800
Subject: [PATCH] feat(mge/uniform): add the judgment of low and high

GitOrigin-RevId: 21788d199211003ce318606e30ca746cf02bffbc
---
 imperative/python/megengine/random/rng.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/imperative/python/megengine/random/rng.py b/imperative/python/megengine/random/rng.py
index e8878a06d..e95eb5d5d 100644
--- a/imperative/python/megengine/random/rng.py
+++ b/imperative/python/megengine/random/rng.py
@@ -106,6 +106,8 @@ def _uniform(
     _ref = Tensor([], dtype="int32", device=device)
     shape = utils.astensor1d(size, _ref, dtype="int32", device=device)
     (output,) = apply(op, shape)
+    if low == 0 and high == 1:
+        return output
     return low + (high - low) * output
 
 
-- 
GitLab