zhuzilin commented on a change in pull request #8041:
URL: https://github.com/apache/tvm/pull/8041#discussion_r634130654
##########
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:
@altanh Sorry that I'm not familiar with the threefry algorithm. Is it
possible to call `_threefry` twice in `threefry_generate` in the following
form? something like:
```python
out_array = irb.buffer_ptr(out_array_ptr)
# deal with most of the array
_threefry(irb, tmp, 0, tmp, 4, out_array, 0, out_len // 4)
if out_len % 4 != 0:
# generate remainders in a small tmp buffer
tmp_array = irb.allocate(gen.dtype, 4, name="tmp", scope="global")
# may need to update the tmp key in between
# ...
_threefry(irb, tmp, 0, tmp, 4, tmp_array, 0, out_len // 4)
# only copy the tmp buffer
for i in range(out_len // 4 * 4, out_len):
out_array[i] = tmp_array[i%4]
```
In this way, we coud avoid copying the whole generated tensor.
--
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]
