FrozenGene commented on a change in pull request #8041:
URL: https://github.com/apache/tvm/pull/8041#discussion_r632401575
##########
File path: tests/python/topi/python/test_topi_prng.py
##########
@@ -118,7 +132,22 @@ def test_threefry_wrapping(target, dev):
), f"{target} does not suppport wrapping unsigned integer arithmetic"
[email protected]_targets
+def test_uniform(target, dev):
+ gen = tvm.relay.random.threefry_key(0).data.asnumpy()
+ m = 1024
+ n = 1024
+ dtype = "float32"
Review comment:
ditto
##########
File path: python/tvm/topi/random/kernel.py
##########
@@ -466,3 +466,54 @@ 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]
Review comment:
What is the meaning of `10`?
##########
File path: 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"
Review comment:
Should cover more types. For example `float64` you have implemented
##########
File path: python/tvm/topi/random/kernel.py
##########
@@ -466,3 +466,54 @@ 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.
Review comment:
What is the reason of `product must be a multiple of 4`?
##########
File path: python/tvm/topi/random/kernel.py
##########
@@ -466,3 +466,54 @@ 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
+
Review comment:
What will be happened if our `out_dtype` are other types? For example
`float16`? Even `int8` / `int32`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
