tkonolige commented on a change in pull request #8041:
URL: https://github.com/apache/tvm/pull/8041#discussion_r633744201
##########
File path: python/tvm/relay/op/random/kernel.py
##########
@@ -132,3 +132,55 @@ def foo(key):
:py:func:`threefry_generate`.
"""
return _make.threefry_split(key)
+
+
+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)
+ (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)
Review comment:
```suggestion
key = threefry_key(0)
key, random_values = uniform(key, (100,), low=0, high=10)
```
##########
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:
I would add a check here anyways just to be safe.
--
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]
