altanh commented on a change in pull request #8041:
URL: https://github.com/apache/tvm/pull/8041#discussion_r633750359
##########
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:
do you mind sending a PR for updating the threefry_generate output, or
rather what approach do you have in mind? I tried to avoid this problem by
truncating output buffer but this required an extra copy, wonder if you have
something else
--
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]
