From 7baf7264c12b1fa88c9344d0a6701fa61573264c Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Fri, 14 May 2021 11:38:59 +0800
Subject: [PATCH 1/9] Add uniform distribution generator wrt threefry PRNG

---
 include/tvm/relay/attrs/random.h           | 11 +++++
 python/tvm/relay/op/op_attrs.py            |  5 ++
 python/tvm/relay/op/random/_kernel.py      |  4 ++
 python/tvm/relay/op/random/kernel.py       | 51 +++++++++++++++++++-
 python/tvm/relay/op/strategy/generic.py    | 23 +++++++++
 python/tvm/topi/random/kernel.py           | 54 ++++++++++++++++++++++
 src/relay/op/random/kernel.cc              | 41 ++++++++++++++++
 tests/python/relay/test_prng.py            | 13 ++++++
 tests/python/topi/python/test_topi_prng.py | 33 +++++++++++++
 9 files changed, 234 insertions(+), 1 deletion(-)

diff --git a/include/tvm/relay/attrs/random.h b/include/tvm/relay/attrs/random.h
index 8238f102dab8..b0ce682f0d54 100644
--- a/include/tvm/relay/attrs/random.h
+++ b/include/tvm/relay/attrs/random.h
@@ -37,6 +37,17 @@ struct ThreefryGenerateAttrs : public tvm::AttrsNode<ThreefryGenerateAttrs> {
   }
 };
 
+struct UniformAttrs : public tvm::AttrsNode<UniformAttrs> {
+  Array<Integer> out_shape;
+  DataType out_dtype;
+
+  TVM_DECLARE_ATTRS(UniformAttrs, "relay.attrs.UniformAttrs") {
+    TVM_ATTR_FIELD(out_shape).describe("Shape of random numbers to generate");
+    TVM_ATTR_FIELD(out_dtype).set_default(NullValue<DataType>())
+        .describe("Data type of the generated numbers");
+  }
+};
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_RANDOM_H_
diff --git a/python/tvm/relay/op/op_attrs.py b/python/tvm/relay/op/op_attrs.py
index 4cc6e0f26b91..b96e8b7a66b5 100644
--- a/python/tvm/relay/op/op_attrs.py
+++ b/python/tvm/relay/op/op_attrs.py
@@ -562,3 +562,8 @@ class BatchToSpaceNDAttrs(Attrs):
 @tvm._ffi.register_object("relay.attrs.ThreefryGenerateAttrs")
 class ThreefryGenerateAttrs(Attrs):
     """Attributes used in ThreefryGenerateAttrs operators"""
+
+
+@tvm._ffi.register_object("relay.attrs.UniformAttrs")
+class UniformAttrs(Attrs):
+    """Attributes used in UniformAttrs operators"""
diff --git a/python/tvm/relay/op/random/_kernel.py b/python/tvm/relay/op/random/_kernel.py
index 8be3397008d5..e601a7073cff 100644
--- a/python/tvm/relay/op/random/_kernel.py
+++ b/python/tvm/relay/op/random/_kernel.py
@@ -27,3 +27,7 @@
 register_pattern("random.threefry_generate", OpPattern.OPAQUE)
 register_strategy("random.threefry_split", strategy.threefry_split_strategy)
 register_pattern("random.threefry_split", OpPattern.OPAQUE)
+
+# Distribution
+register_strategy("random.uniform", strategy.uniform_strategy)
+register_pattern("random.uniform", OpPattern.OPAQUE)
diff --git a/python/tvm/relay/op/random/kernel.py b/python/tvm/relay/op/random/kernel.py
index 96634943128d..74ec9a39b122 100644
--- a/python/tvm/relay/op/random/kernel.py
+++ b/python/tvm/relay/op/random/kernel.py
@@ -21,7 +21,7 @@
 import sys
 import numpy as np
 
-from ...expr import Constant
+from ...expr import Constant, Expr, const
 from .... import nd
 from . import _make
 
@@ -132,3 +132,52 @@ def foo(key):
         :py:func:`threefry_generate`.
     """
     return _make.threefry_split(key)
+
+
+def uniform(key, shape, dtype="float32", low=0., high=1.):
+    """Draw samples from a uniform distribution.
+
+    Samples are uniformly distributed over the half-open interval [low, high)
+    (includes low, but excludes high). In other words, any value within the
+    given interval is equally likely to be drawn by uniform.
+
+    Example
+    -------
+
+    .. code-block:: python
+
+        key = threefry_key(0)
+        random_values = uniform(key, (100,), low=0, high=10)
+
+    Parameters
+    ----------
+    key : relay.Expr
+        key that uniquely determines the random values. Multiple uses with the
+        same generator will generate the same random values. This generator should be
+        treated as an opaque pointer. You can create one from calling
+        :py:func:`threefry_key`, :py:func:`threefry_split`, or
+        :py:func:`threefry_generate`. **Do not use this generator again after calling
+        this function.**
+
+    shape : Sequence[int]
+        Desired outputs shape of random numbers.
+
+    dtype : str
+        Desired outputs type of random numbers.
+
+    low : float or relay.Expr, optional
+        Lower bound of the uniform distribution.
+
+    high : float or relay.Expr, optional
+        Upper bound of the uniform distribution.
+
+    Returns
+    -------
+    random_values : relay.Expr
+        The generated uniform distributed random numbers.
+    """
+    if not isinstance(low, Expr):
+        low = const(low, dtype=dtype)
+    if not isinstance(high, Expr):
+        high = const(high, dtype=dtype)
+    return _make.uniform(key, low, high, shape, dtype)
diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py
index 7451b397265f..61de68315780 100644
--- a/python/tvm/relay/op/strategy/generic.py
+++ b/python/tvm/relay/op/strategy/generic.py
@@ -1495,6 +1495,29 @@ def threefry_split_strategy(attrs, inputs, out_type, target):
     return strategy
 
 
+# uniform
+def wrap_compute_uniform(topi_compute):
+    """Wrap uniform topi compute"""
+
+    def _compute_uniform(attrs, inputs, _):
+        return [topi_compute(inputs[0], inputs[1], inputs[2],
+                             attrs.out_shape, attrs.out_dtype)]
+
+    return _compute_uniform
+
+
+@override_native_generic_func("uniform_strategy")
+def uniform_strategy(attrs, inputs, out_type, target):
+    """uniform generic strategy"""
+    strategy = _op.OpStrategy()
+    strategy.add_implementation(
+        wrap_compute_uniform(topi.random.uniform),
+        wrap_topi_schedule(topi.generic.schedule_extern),
+        name="uniform.generic",
+    )
+    return strategy
+
+
 def wrap_compute_scanop(topi_compute):
     """Wrap scanop style topi compute"""
 
diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
index b6c0b3fa5930..261d1ac8f79b 100644
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -466,3 +466,57 @@ def gen_ir(out_ptr):
     out_ary = tvm.nd.array(np.ones((1,), "uint64"), device)
     tvm.build(s, [f], target=target)(out_ary)
     return out_ary.asnumpy()[0] == 0
+
+
+def uniform(gen, low, high, out_shape, out_dtype):
+    """Draw samples from a uniform distribution.
+
+    Samples are uniformly distributed over the half-open interval [low, high)
+    (includes low, but excludes high). In other words, any value within the
+    given interval is equally likely to be drawn by uniform.
+
+    Parameters
+    ----------
+    gen : Tensor[10, uint64]
+        Generator state. Can be create with :py:func:`tvm.relay.threefry_key`. This should not be
+        reused in another function, otherwise random numbers will be repeated.
+
+    low : Tensor[(), out_dtype]
+        Lower boundary of the output interval. All values generated will be
+        greater than or equal to low.
+
+    high : Tensor[(), out_dtype]
+        Upper boundary of the output interval. All values generated will be
+        less than high.
+
+    out_shape : Sequence[int]
+        Output shape of the random numbers. Product of all dimensions must be a multiple of 4.
+
+    out_dtype : str
+        The output dtype.
+
+    Returns
+    -------
+    out : Tensor[out_shape, out_dtype]
+        Tensor of random numbers with shape `out_shape` and type `out_dtype`.
+    """
+    _, random_bits = threefry_generate(gen, out_shape)
+    nbits = 64
+    if out_dtype == "float32":
+        nfraction = 23
+    elif out_dtype == "float64":
+        nfraction = 52
+
+    def uniform_scalar(bits):
+        bits = bits >> (nbits - nfraction)
+        standard_uniform = bits.astype(out_dtype) / float(1 << nfraction)
+        return standard_uniform
+
+    standard_uniform_values = tvm.te.compute(
+        out_shape, lambda *i : uniform_scalar(random_bits(*i)))
+
+    uniform_values = tvm.topi.add(
+                        tvm.topi.multiply(standard_uniform_values, high - low),
+                        low)
+
+    return uniform_values
diff --git a/src/relay/op/random/kernel.cc b/src/relay/op/random/kernel.cc
index ec092a7e05f2..744eb52dd0a5 100644
--- a/src/relay/op/random/kernel.cc
+++ b/src/relay/op/random/kernel.cc
@@ -85,5 +85,46 @@ RELAY_REGISTER_OP("random.threefry_split")
     .add_argument("key", "Tensor", "Input Threefry key")
     .add_type_rel("ThreefrySplit", ThreefrySplitRel);
 
+TVM_REGISTER_NODE_TYPE(UniformAttrs);
+
+bool UniformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                const TypeReporter& reporter) {
+  const UniformAttrs* param = attrs.as<UniformAttrs>();
+  ICHECK_EQ(types.size(), 4) << "ThreefryGenerate should have one input and one output";
+
+  std::vector<IndexExpr> oshape;
+  for (auto& x : param->out_shape) {
+    oshape.push_back(x);
+  }
+  DataType out_dtype = param->out_dtype;
+
+  reporter->Assign(types[0], ThreefryKeyType());
+  reporter->Assign(types[1], TensorType({}, out_dtype));
+  reporter->Assign(types[2], TensorType({}, out_dtype));
+  // generate returns the next key and an array of random values
+  reporter->Assign(types[3], TensorType(oshape, out_dtype));
+  return true;
+}
+
+Expr MakeUniform(Expr key, Expr low, Expr high, Array<Integer> out_shape, DataType out_dtype) {
+  auto attrs = make_object<UniformAttrs>();
+  attrs->out_shape = out_shape;
+  attrs->out_dtype = out_dtype;
+  static const Op& op = Op::Get("random.uniform");
+  return Call(op, {key, low, high}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op.random._make.uniform").set_body_typed(MakeUniform);
+
+RELAY_REGISTER_OP("random.uniform")
+    .describe(
+        R"doc(Generate an array of random numbers under uniform distribution.)doc" TVM_ADD_FILELINE)
+    .set_num_inputs(3)
+    .set_attrs_type<UniformAttrs>()
+    .add_argument("key", "Tensor", "Input Threefry key")
+    .add_argument("low", "Tensor", "Lower bound of the distribution")
+    .add_argument("high", "Tensor", "Higher bound of the distribution")
+    .add_type_rel("Uniform", UniformRel);
+
 }  // namespace relay
 }  // namespace tvm
diff --git a/tests/python/relay/test_prng.py b/tests/python/relay/test_prng.py
index ba4fdc466ecc..51938cd66330 100644
--- a/tests/python/relay/test_prng.py
+++ b/tests/python/relay/test_prng.py
@@ -103,6 +103,19 @@ def test_threefry_split_infer():
     assert tvm.ir.structural_equal(f.ret_type, expected_type)
 
 
+def test_uniform_infer():
+    oshape = (12,)
+    odtype = "float32"
+    gen_type = tvm.relay.TensorType(oshape, dtype=odtype)
+    expected_type = gen_type
+
+    key = tvm.relay.random.threefry_key(1)
+    out_keys = tvm.relay.random.uniform(key, oshape, odtype)
+    f = tvm.relay.Function([], out_keys)
+    f = run_infer_type(f)
+    assert tvm.ir.structural_equal(f.ret_type, expected_type)
+
+
 @pytest.mark.xfail(raises=tvm.error.TVMError)
 def test_threefry_generate_infer_fail():
     # xfail: key size should be 10
diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py
index 4ad3a80c6a9e..1a7cda782025 100644
--- a/tests/python/topi/python/test_topi_prng.py
+++ b/tests/python/topi/python/test_topi_prng.py
@@ -43,6 +43,23 @@ def threefry_generate(target, dev, gen, size):
     return out_gen.asnumpy(), rands.asnumpy()
 
 
+def uniform(target, dev, gen, low, high, size, dtype):
+    gen_placeholder = tvm.te.placeholder(gen.shape, name="gen", dtype="uint64")
+    low_placeholder = tvm.te.placeholder(low.shape, name="low", dtype=dtype)
+    high_placeholder = tvm.te.placeholder(high.shape, name="high", dtype=dtype)
+    print(low_placeholder)
+    print(high_placeholder)
+    out_placeholder = tvm.topi.random.uniform(gen_placeholder, low_placeholder,
+                                              high_placeholder, size, dtype)
+    print(out_placeholder)
+    s = tvm.topi.generic.schedule_extern([out_placeholder])
+    f = tvm.build(s, [gen_placeholder, low_placeholder, high_placeholder,
+                      out_placeholder])
+    rands = tvm.nd.array(np.zeros(size, dtype=dtype))
+    f(tvm.nd.array(gen), tvm.nd.array(low), tvm.nd.array(high), rands)
+    return rands.asnumpy()
+
+
 @tvm.testing.parametrize_targets
 def test_threefry_split(target, dev):
     # test that results of split do not equal eachother or the input
@@ -118,7 +135,23 @@ def test_threefry_wrapping(target, dev):
     ), f"{target} does not suppport wrapping unsigned integer arithmetic"
 
 
+@tvm.testing.parametrize_targets
+def test_uniform(target, dev):
+    gen = tvm.relay.random.threefry_key(0).data.asnumpy()
+    m = 1024
+    n = 1024
+    dtype = "float32"
+    low = np.array(5.0, dtype=dtype)
+    high = np.array(10.0, dtype=dtype)
+    rands = uniform(
+        target, dev, gen, low, high, (m, n), "float32")
+    assert abs(np.mean(rands) - 7.5) < 1e-1
+    assert abs(np.min(rands) - 5.0) < 1e-3
+    assert abs(np.max(rands) - 10.0) < 1e-3
+
+
 if __name__ == "__main__":
     test_threefry_split(tvm.target.Target("llvm"), tvm.device("cpu"))
     test_threefry_generate(tvm.target.Target("llvm"), tvm.device("cpu"))
     test_threefry_wrapping(tvm.target.Target("llvm"), tvm.device("cpu"))
+    test_uniform(tvm.target.Target("llvm"), tvm.device("cpu"))

From fd68ff48a50abfae09885b5b8f91762a6d9b10e1 Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Fri, 14 May 2021 12:23:01 +0800
Subject: [PATCH 2/9] fix lint

---
 include/tvm/relay/attrs/random.h           |  3 ++-
 python/tvm/relay/op/random/kernel.py       |  2 +-
 python/tvm/relay/op/strategy/generic.py    |  3 +--
 python/tvm/topi/random/kernel.py           |  7 ++-----
 tests/python/topi/python/test_topi_prng.py | 11 +++++------
 5 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/include/tvm/relay/attrs/random.h b/include/tvm/relay/attrs/random.h
index b0ce682f0d54..46cab8831caf 100644
--- a/include/tvm/relay/attrs/random.h
+++ b/include/tvm/relay/attrs/random.h
@@ -43,7 +43,8 @@ struct UniformAttrs : public tvm::AttrsNode<UniformAttrs> {
 
   TVM_DECLARE_ATTRS(UniformAttrs, "relay.attrs.UniformAttrs") {
     TVM_ATTR_FIELD(out_shape).describe("Shape of random numbers to generate");
-    TVM_ATTR_FIELD(out_dtype).set_default(NullValue<DataType>())
+    TVM_ATTR_FIELD(out_dtype)
+        .set_default(NullValue<DataType>())
         .describe("Data type of the generated numbers");
   }
 };
diff --git a/python/tvm/relay/op/random/kernel.py b/python/tvm/relay/op/random/kernel.py
index 74ec9a39b122..f728eb6e9da5 100644
--- a/python/tvm/relay/op/random/kernel.py
+++ b/python/tvm/relay/op/random/kernel.py
@@ -134,7 +134,7 @@ def foo(key):
     return _make.threefry_split(key)
 
 
-def uniform(key, shape, dtype="float32", low=0., high=1.):
+def uniform(key, shape, dtype="float32", low=0.0, high=1.0):
     """Draw samples from a uniform distribution.
 
     Samples are uniformly distributed over the half-open interval [low, high)
diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py
index 61de68315780..f7747a9db703 100644
--- a/python/tvm/relay/op/strategy/generic.py
+++ b/python/tvm/relay/op/strategy/generic.py
@@ -1500,8 +1500,7 @@ def wrap_compute_uniform(topi_compute):
     """Wrap uniform topi compute"""
 
     def _compute_uniform(attrs, inputs, _):
-        return [topi_compute(inputs[0], inputs[1], inputs[2],
-                             attrs.out_shape, attrs.out_dtype)]
+        return [topi_compute(inputs[0], inputs[1], inputs[2], attrs.out_shape, attrs.out_dtype)]
 
     return _compute_uniform
 
diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
index 261d1ac8f79b..6e48c81e9b64 100644
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -512,11 +512,8 @@ def uniform_scalar(bits):
         standard_uniform = bits.astype(out_dtype) / float(1 << nfraction)
         return standard_uniform
 
-    standard_uniform_values = tvm.te.compute(
-        out_shape, lambda *i : uniform_scalar(random_bits(*i)))
+    standard_uniform_values = tvm.te.compute(out_shape, lambda *i: uniform_scalar(random_bits(*i)))
 
-    uniform_values = tvm.topi.add(
-                        tvm.topi.multiply(standard_uniform_values, high - low),
-                        low)
+    uniform_values = tvm.topi.add(tvm.topi.multiply(standard_uniform_values, high - low), low)
 
     return uniform_values
diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py
index 1a7cda782025..89b6b9e59bfd 100644
--- a/tests/python/topi/python/test_topi_prng.py
+++ b/tests/python/topi/python/test_topi_prng.py
@@ -49,12 +49,12 @@ def uniform(target, dev, gen, low, high, size, dtype):
     high_placeholder = tvm.te.placeholder(high.shape, name="high", dtype=dtype)
     print(low_placeholder)
     print(high_placeholder)
-    out_placeholder = tvm.topi.random.uniform(gen_placeholder, low_placeholder,
-                                              high_placeholder, size, dtype)
+    out_placeholder = tvm.topi.random.uniform(
+        gen_placeholder, low_placeholder, high_placeholder, size, dtype
+    )
     print(out_placeholder)
     s = tvm.topi.generic.schedule_extern([out_placeholder])
-    f = tvm.build(s, [gen_placeholder, low_placeholder, high_placeholder,
-                      out_placeholder])
+    f = tvm.build(s, [gen_placeholder, low_placeholder, high_placeholder, out_placeholder])
     rands = tvm.nd.array(np.zeros(size, dtype=dtype))
     f(tvm.nd.array(gen), tvm.nd.array(low), tvm.nd.array(high), rands)
     return rands.asnumpy()
@@ -143,8 +143,7 @@ def test_uniform(target, dev):
     dtype = "float32"
     low = np.array(5.0, dtype=dtype)
     high = np.array(10.0, dtype=dtype)
-    rands = uniform(
-        target, dev, gen, low, high, (m, n), "float32")
+    rands = uniform(target, dev, gen, low, high, (m, n), "float32")
     assert abs(np.mean(rands) - 7.5) < 1e-1
     assert abs(np.min(rands) - 5.0) < 1e-3
     assert abs(np.max(rands) - 10.0) < 1e-3

From 135da627b78408e96dab9249c6c3c011d4ca54a6 Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Fri, 14 May 2021 13:01:57 +0800
Subject: [PATCH 3/9] remove the redundant print

---
 tests/python/topi/python/test_topi_prng.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py
index 89b6b9e59bfd..fc7e134d4c61 100644
--- a/tests/python/topi/python/test_topi_prng.py
+++ b/tests/python/topi/python/test_topi_prng.py
@@ -47,12 +47,9 @@ def uniform(target, dev, gen, low, high, size, dtype):
     gen_placeholder = tvm.te.placeholder(gen.shape, name="gen", dtype="uint64")
     low_placeholder = tvm.te.placeholder(low.shape, name="low", dtype=dtype)
     high_placeholder = tvm.te.placeholder(high.shape, name="high", dtype=dtype)
-    print(low_placeholder)
-    print(high_placeholder)
     out_placeholder = tvm.topi.random.uniform(
         gen_placeholder, low_placeholder, high_placeholder, size, dtype
     )
-    print(out_placeholder)
     s = tvm.topi.generic.schedule_extern([out_placeholder])
     f = tvm.build(s, [gen_placeholder, low_placeholder, high_placeholder, out_placeholder])
     rands = tvm.nd.array(np.zeros(size, dtype=dtype))

From cbaec73e56f6f58dbf56058ab2b3122707be1b18 Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Sun, 16 May 2021 00:09:34 +0800
Subject: [PATCH 4/9] modifications based on review

---
 python/tvm/relay/op/strategy/generic.py    |  2 +-
 python/tvm/topi/random/kernel.py           |  8 +++---
 src/relay/op/random/kernel.cc              |  4 +--
 tests/python/relay/test_prng.py            | 20 ++++++++-------
 tests/python/topi/python/test_topi_prng.py | 29 +++++++++++++---------
 5 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py
index f7747a9db703..97d1a91f744f 100644
--- a/python/tvm/relay/op/strategy/generic.py
+++ b/python/tvm/relay/op/strategy/generic.py
@@ -1500,7 +1500,7 @@ def wrap_compute_uniform(topi_compute):
     """Wrap uniform topi compute"""
 
     def _compute_uniform(attrs, inputs, _):
-        return [topi_compute(inputs[0], inputs[1], inputs[2], attrs.out_shape, attrs.out_dtype)]
+        return list(topi_compute(inputs[0], inputs[1], inputs[2], attrs.out_shape, attrs.out_dtype))
 
     return _compute_uniform
 
diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
index 6e48c81e9b64..2d93afe2ebe7 100644
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -477,7 +477,7 @@ def uniform(gen, low, high, out_shape, out_dtype):
 
     Parameters
     ----------
-    gen : Tensor[10, uint64]
+    gen : ThreefryKey
         Generator state. Can be create with :py:func:`tvm.relay.threefry_key`. This should not be
         reused in another function, otherwise random numbers will be repeated.
 
@@ -500,7 +500,7 @@ def uniform(gen, low, high, out_shape, out_dtype):
     out : Tensor[out_shape, out_dtype]
         Tensor of random numbers with shape `out_shape` and type `out_dtype`.
     """
-    _, random_bits = threefry_generate(gen, out_shape)
+    new_gen, random_bits = threefry_generate(gen, out_shape)
     nbits = 64
     if out_dtype == "float32":
         nfraction = 23
@@ -509,11 +509,11 @@ def uniform(gen, low, high, out_shape, out_dtype):
 
     def uniform_scalar(bits):
         bits = bits >> (nbits - nfraction)
-        standard_uniform = bits.astype(out_dtype) / float(1 << nfraction)
+        standard_uniform = bits.astype(out_dtype) / tir.const(1 << nfraction, dtype=out_dtype)
         return standard_uniform
 
     standard_uniform_values = tvm.te.compute(out_shape, lambda *i: uniform_scalar(random_bits(*i)))
 
     uniform_values = tvm.topi.add(tvm.topi.multiply(standard_uniform_values, high - low), low)
 
-    return uniform_values
+    return new_gen, uniform_values
diff --git a/src/relay/op/random/kernel.cc b/src/relay/op/random/kernel.cc
index 744eb52dd0a5..b2513264d0ba 100644
--- a/src/relay/op/random/kernel.cc
+++ b/src/relay/op/random/kernel.cc
@@ -90,7 +90,7 @@ TVM_REGISTER_NODE_TYPE(UniformAttrs);
 bool UniformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
                 const TypeReporter& reporter) {
   const UniformAttrs* param = attrs.as<UniformAttrs>();
-  ICHECK_EQ(types.size(), 4) << "ThreefryGenerate should have one input and one output";
+  ICHECK_EQ(types.size(), 4) << "Uniform should have three input and one output";
 
   std::vector<IndexExpr> oshape;
   for (auto& x : param->out_shape) {
@@ -102,7 +102,7 @@ bool UniformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   reporter->Assign(types[1], TensorType({}, out_dtype));
   reporter->Assign(types[2], TensorType({}, out_dtype));
   // generate returns the next key and an array of random values
-  reporter->Assign(types[3], TensorType(oshape, out_dtype));
+  reporter->Assign(types[3], TupleType({ThreefryKeyType(), TensorType(oshape, out_dtype)}));
   return true;
 }
 
diff --git a/tests/python/relay/test_prng.py b/tests/python/relay/test_prng.py
index 51938cd66330..f622bfca5625 100644
--- a/tests/python/relay/test_prng.py
+++ b/tests/python/relay/test_prng.py
@@ -105,15 +105,17 @@ def test_threefry_split_infer():
 
 def test_uniform_infer():
     oshape = (12,)
-    odtype = "float32"
-    gen_type = tvm.relay.TensorType(oshape, dtype=odtype)
-    expected_type = gen_type
-
-    key = tvm.relay.random.threefry_key(1)
-    out_keys = tvm.relay.random.uniform(key, oshape, odtype)
-    f = tvm.relay.Function([], out_keys)
-    f = run_infer_type(f)
-    assert tvm.ir.structural_equal(f.ret_type, expected_type)
+    odtypes = ["float32", "float64"]
+    for odtype in odtypes:
+        key_type = tvm.relay.TensorType([10], dtype="uint64")
+        gen_type = tvm.relay.TensorType(oshape, dtype=odtype)
+        expected_type = tvm.relay.TupleType([key_type, gen_type])
+
+        key = tvm.relay.random.threefry_key(1)
+        rand1 = tvm.relay.random.uniform(key, oshape, odtype)
+        f = tvm.relay.Function([], rand1)
+        f = run_infer_type(f)
+        assert tvm.ir.structural_equal(f.ret_type, expected_type)
 
 
 @pytest.mark.xfail(raises=tvm.error.TVMError)
diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py
index fc7e134d4c61..2d773cc50588 100644
--- a/tests/python/topi/python/test_topi_prng.py
+++ b/tests/python/topi/python/test_topi_prng.py
@@ -47,14 +47,17 @@ def uniform(target, dev, gen, low, high, size, dtype):
     gen_placeholder = tvm.te.placeholder(gen.shape, name="gen", dtype="uint64")
     low_placeholder = tvm.te.placeholder(low.shape, name="low", dtype=dtype)
     high_placeholder = tvm.te.placeholder(high.shape, name="high", dtype=dtype)
-    out_placeholder = tvm.topi.random.uniform(
+    left_placeholder, right_placeholder = tvm.topi.random.uniform(
         gen_placeholder, low_placeholder, high_placeholder, size, dtype
     )
-    s = tvm.topi.generic.schedule_extern([out_placeholder])
-    f = tvm.build(s, [gen_placeholder, low_placeholder, high_placeholder, out_placeholder])
+    s = tvm.topi.generic.schedule_extern([left_placeholder, right_placeholder])
+    f = tvm.build(
+        s, [gen_placeholder, low_placeholder, high_placeholder, left_placeholder, right_placeholder]
+    )
+    out_gen = tvm.nd.array(np.zeros(gen.shape, dtype="uint64"))
     rands = tvm.nd.array(np.zeros(size, dtype=dtype))
-    f(tvm.nd.array(gen), tvm.nd.array(low), tvm.nd.array(high), rands)
-    return rands.asnumpy()
+    f(tvm.nd.array(gen), tvm.nd.array(low), tvm.nd.array(high), out_gen, rands)
+    return out_gen.asnumpy(), rands.asnumpy()
 
 
 @tvm.testing.parametrize_targets
@@ -137,13 +140,15 @@ def test_uniform(target, dev):
     gen = tvm.relay.random.threefry_key(0).data.asnumpy()
     m = 1024
     n = 1024
-    dtype = "float32"
-    low = np.array(5.0, dtype=dtype)
-    high = np.array(10.0, dtype=dtype)
-    rands = uniform(target, dev, gen, low, high, (m, n), "float32")
-    assert abs(np.mean(rands) - 7.5) < 1e-1
-    assert abs(np.min(rands) - 5.0) < 1e-3
-    assert abs(np.max(rands) - 10.0) < 1e-3
+    dtypes = ["float32", "float64"]
+    for dtype in dtypes:
+        low = np.array(5.0, dtype=dtype)
+        high = np.array(10.0, dtype=dtype)
+        new_gen, rands = uniform(target, dev, gen, low, high, (m, n), dtype)
+        assert (gen != new_gen).any()
+        assert abs(np.mean(rands) - 7.5) < 1e-1
+        assert abs(np.min(rands) - 5.0) < 1e-3
+        assert abs(np.max(rands) - 10.0) < 1e-3
 
 
 if __name__ == "__main__":

From 024e1765b28eefb0cfcc24a6bce97c128ac25981 Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Sun, 16 May 2021 00:12:45 +0800
Subject: [PATCH 5/9] update docs

---
 python/tvm/relay/op/random/kernel.py | 3 +++
 python/tvm/topi/random/kernel.py     | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/python/tvm/relay/op/random/kernel.py b/python/tvm/relay/op/random/kernel.py
index f728eb6e9da5..415c19134c02 100644
--- a/python/tvm/relay/op/random/kernel.py
+++ b/python/tvm/relay/op/random/kernel.py
@@ -173,6 +173,9 @@ def uniform(key, shape, dtype="float32", low=0.0, high=1.0):
 
     Returns
     -------
+    new_key : relay.Expr
+        New random key to pass to future uses of random functions.
+
     random_values : relay.Expr
         The generated uniform distributed random numbers.
     """
diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
index 2d93afe2ebe7..f4fe2f9195e4 100644
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -497,6 +497,9 @@ def uniform(gen, low, high, out_shape, out_dtype):
 
     Returns
     -------
+    new_gen : ThreefryKey
+        New generator state that is distinct from `gen`.
+
     out : Tensor[out_shape, out_dtype]
         Tensor of random numbers with shape `out_shape` and type `out_dtype`.
     """

From fd1368bdf6d99f72f5b0f7cba9ca9ba3a6b45e48 Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Sun, 16 May 2021 10:31:16 +0800
Subject: [PATCH 6/9] update uniform algorithm to use bit operations only

---
 python/tvm/topi/random/kernel.py | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
index f4fe2f9195e4..308064bed623 100644
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -504,19 +504,28 @@ def uniform(gen, low, high, out_shape, out_dtype):
         Tensor of random numbers with shape `out_shape` and type `out_dtype`.
     """
     new_gen, random_bits = threefry_generate(gen, out_shape)
-    nbits = 64
     if out_dtype == "float32":
+        random_dtype = "uint32"
+        nbits = 32
         nfraction = 23
     elif out_dtype == "float64":
+        random_dtype = "uint64"
+        nbits = 64
         nfraction = 52
+    nexp = nbits - nfraction - 1
+    random_bits = random_bits.astype(random_dtype)
 
-    def uniform_scalar(bits):
-        bits = bits >> (nbits - nfraction)
-        standard_uniform = bits.astype(out_dtype) / tir.const(1 << nfraction, dtype=out_dtype)
-        return standard_uniform
-
-    standard_uniform_values = tvm.te.compute(out_shape, lambda *i: uniform_scalar(random_bits(*i)))
-
+    fraction = tvm.topi.right_shift(
+        random_bits, tvm.tir.const(nbits - nfraction, dtype=random_dtype)
+    )
+    exponent = tvm.topi.left_shift(
+        tvm.topi.full(out_shape, random_dtype, (1 << (nexp - 1)) - 1),
+        tvm.tir.const(nfraction, dtype=random_dtype),
+    )
+    mantissa = tvm.topi.bitwise_or(fraction, exponent).astype(random_dtype)
+    standard_uniform_values = tvm.topi.reinterpret(mantissa, out_dtype) - tvm.tir.const(
+        1, dtype=out_dtype
+    )
     uniform_values = tvm.topi.add(tvm.topi.multiply(standard_uniform_values, high - low), low)
 
     return new_gen, uniform_values

From bf577e4989503eef5a0ad53658d635a13ebe1fa9 Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Mon, 17 May 2021 23:08:29 +0800
Subject: [PATCH 7/9] add type restrictions

---
 src/relay/op/random/kernel.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/relay/op/random/kernel.cc b/src/relay/op/random/kernel.cc
index b2513264d0ba..0467ac8413fb 100644
--- a/src/relay/op/random/kernel.cc
+++ b/src/relay/op/random/kernel.cc
@@ -97,7 +97,8 @@ bool UniformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
     oshape.push_back(x);
   }
   DataType out_dtype = param->out_dtype;
-
+  // we are supporting float32 and float64 at the moment.
+  ICHECK(out_dtype.is_float() && (out_dtype.bits() == 32 || out_dtype.bits() == 64));
   reporter->Assign(types[0], ThreefryKeyType());
   reporter->Assign(types[1], TensorType({}, out_dtype));
   reporter->Assign(types[2], TensorType({}, out_dtype));

From aa372d4422c8ac7d3836ace5c59b1979447b396f Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Tue, 18 May 2021 16:18:50 +0800
Subject: [PATCH 8/9] minor fix upon review

---
 python/tvm/relay/op/random/kernel.py | 2 +-
 python/tvm/topi/random/kernel.py     | 1 +
 src/relay/op/random/kernel.cc        | 7 ++++++-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relay/op/random/kernel.py b/python/tvm/relay/op/random/kernel.py
index 415c19134c02..fc1248e85678 100644
--- a/python/tvm/relay/op/random/kernel.py
+++ b/python/tvm/relay/op/random/kernel.py
@@ -147,7 +147,7 @@ def uniform(key, shape, dtype="float32", low=0.0, high=1.0):
     .. code-block:: python
 
         key = threefry_key(0)
-        random_values = uniform(key, (100,), low=0, high=10)
+        key, random_values = uniform(key, (100,), low=0, high=10)
 
     Parameters
     ----------
diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
index 308064bed623..83e88ed096f1 100644
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -504,6 +504,7 @@ def uniform(gen, low, high, out_shape, out_dtype):
         Tensor of random numbers with shape `out_shape` and type `out_dtype`.
     """
     new_gen, random_bits = threefry_generate(gen, out_shape)
+    assert out_dtype in ("float32", "float64")
     if out_dtype == "float32":
         random_dtype = "uint32"
         nbits = 32
diff --git a/src/relay/op/random/kernel.cc b/src/relay/op/random/kernel.cc
index 0467ac8413fb..30ceff13c5ed 100644
--- a/src/relay/op/random/kernel.cc
+++ b/src/relay/op/random/kernel.cc
@@ -98,7 +98,12 @@ bool UniformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   }
   DataType out_dtype = param->out_dtype;
   // we are supporting float32 and float64 at the moment.
-  ICHECK(out_dtype.is_float() && (out_dtype.bits() == 32 || out_dtype.bits() == 64));
+  if (!(out_dtype.is_float() && (out_dtype.bits() == 32 || out_dtype.bits() == 64))) {
+    reporter->GetDiagCtx().EmitFatal(Diagnostic::Error(reporter->GetSpan())
+                                     << "We only support generating uniform random value of "
+                                     << "type float32 or float64, got " << out_dtype << ".");
+    return false;
+  }
   reporter->Assign(types[0], ThreefryKeyType());
   reporter->Assign(types[1], TensorType({}, out_dtype));
   reporter->Assign(types[2], TensorType({}, out_dtype));

From c707ce0e884998953104e6fd4de03ef1e282164d Mon Sep 17 00:00:00 2001
From: zilinzhu <zilinzhu@tencent.com>
Date: Wed, 19 May 2021 11:45:10 +0800
Subject: [PATCH 9/9] update test and error information

---
 python/tvm/topi/random/kernel.py           | 4 +++-
 src/relay/op/random/kernel.cc              | 2 +-
 tests/python/topi/python/test_topi_prng.py | 4 ++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
index 83e88ed096f1..2831baba30cd 100644
--- a/python/tvm/topi/random/kernel.py
+++ b/python/tvm/topi/random/kernel.py
@@ -504,7 +504,9 @@ def uniform(gen, low, high, out_shape, out_dtype):
         Tensor of random numbers with shape `out_shape` and type `out_dtype`.
     """
     new_gen, random_bits = threefry_generate(gen, out_shape)
-    assert out_dtype in ("float32", "float64")
+    assert out_dtype in ("float32", "float64"), (
+        "Only support float32 or float64 for now, got %s" % out_dtype
+    )
     if out_dtype == "float32":
         random_dtype = "uint32"
         nbits = 32
diff --git a/src/relay/op/random/kernel.cc b/src/relay/op/random/kernel.cc
index 30ceff13c5ed..077a1ea11673 100644
--- a/src/relay/op/random/kernel.cc
+++ b/src/relay/op/random/kernel.cc
@@ -90,7 +90,7 @@ TVM_REGISTER_NODE_TYPE(UniformAttrs);
 bool UniformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
                 const TypeReporter& reporter) {
   const UniformAttrs* param = attrs.as<UniformAttrs>();
-  ICHECK_EQ(types.size(), 4) << "Uniform should have three input and one output";
+  ICHECK_EQ(types.size(), 4) << "Uniform should have three inputs and one output";
 
   std::vector<IndexExpr> oshape;
   for (auto& x : param->out_shape) {
diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py
index 2d773cc50588..9189fa792898 100644
--- a/tests/python/topi/python/test_topi_prng.py
+++ b/tests/python/topi/python/test_topi_prng.py
@@ -147,8 +147,8 @@ def test_uniform(target, dev):
         new_gen, rands = uniform(target, dev, gen, low, high, (m, n), dtype)
         assert (gen != new_gen).any()
         assert abs(np.mean(rands) - 7.5) < 1e-1
-        assert abs(np.min(rands) - 5.0) < 1e-3
-        assert abs(np.max(rands) - 10.0) < 1e-3
+        assert np.min(rands) >= 5.0
+        assert np.max(rands) <= 10.0
 
 
 if __name__ == "__main__":