tkonolige commented on a change in pull request #8041:
URL: https://github.com/apache/tvm/pull/8041#discussion_r632639897
##########
File path: python/tvm/relay/op/random/kernel.py
##########
@@ -132,3 +132,52 @@ 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)
+
+ 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
Review comment:
One thing we want to think about here is the API for generating random
values. The main question is around handling the threefry key. Do we want the
user to split the key before every invocation, or do we want anything that uses
the key to return the next one.
i.e.
```
key_1, key_2 = split(key)
data = uniform(key_2, ...)
my_other_random_funtion(key_1)
```
vs
```
data, key = uniform(key, ...)
my_other_random_function(key)
```
Personally I prefer the second. It also makes better use of the key space
and avoid recomputing the key too often.
##########
File path: 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";
Review comment:
```suggestion
ICHECK_EQ(types.size(), 4) << "Uniform should have one input and one
output";
```
##########
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
+
+ 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)
Review comment:
How well does this approach work when we have a large range (high -
low)? It seems like we would be loosing a lot of potential randomness.
--
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]
