yongwww commented on code in PR #16575: URL: https://github.com/apache/tvm/pull/16575#discussion_r1491938967
########## python/tvm/relax/frontend/nn/op.py: ########## @@ -1825,3 +1826,260 @@ def print_(tensor: Tensor): filename, line_number = inspect.getframeinfo(inspect.currentframe().f_back)[:2] line_info = f"{filename}:{line_number}" debug_func("vm.builtin.debug_print", tensor, _line_info=line_info) + + +def less(a: Tensor, b: Tensor, name: str = "less") -> Tensor: + """Broadcasted element-wise comparison for (lhs < rhs). + + Parameters + ---------- + a : Tensor + The first input tensor. + + b : Tensor + The second input tensor. + + name : str + Name hint. + + Returns + ------- + result : Tensor + The computed result. + """ + return wrap_nested(_op.less(a._expr, b._expr), name) + + +def less_equal(a: Tensor, b: Tensor, name: str = "less_equal") -> Tensor: + """Broadcasted element-wise comparison for (lhs <= rhs). + + Parameters + ---------- + a : Tensor + The first input tensor. + + b : Tensor + The second input tensor. + + name : str + Name hint. + + Returns + ------- + result : Tensor + The computed result. + """ + return wrap_nested(_op.less_equal(a._expr, b._expr), name) + + +def greater(a: Tensor, b: Tensor, name: str = "greater") -> Tensor: + """Broadcasted element-wise comparison for (lhs > rhs). + + Parameters + ---------- + a : Tensor + The first input tensor. + + b : Tensor + The second input tensor. + + name : str + Name hint. + + Returns + ------- + result : Tensor + The computed result. + """ + return wrap_nested(_op.greater(a._expr, b._expr), name) + + +def greater_equal(a: Tensor, b: Tensor, name: str = "greater_equal") -> Tensor: + """Broadcasted element-wise comparison for (lhs >= rhs). + + Parameters + ---------- + a : Tensor + The first input tensor. + + b : Tensor + The second input tensor. + + name : str + Name hint. + + Returns + ------- + result : Tensor + The computed result. + """ + return wrap_nested(_op.greater_equal(a._expr, b._expr), name) + + +def equal(a: Tensor, b: Tensor, name: str = "equal") -> Tensor: + """Broadcasted element-wise comparison for (lhs == rhs). + + Parameters + ---------- + a : Tensor + The first input tensor. + + b : Tensor + The second input tensor. + + name : str + Name hint. + + Returns + ------- + result : Tensor + The computed result. + """ + return wrap_nested(_op.equal(a._expr, b._expr), name) + + +def not_equal(a: Tensor, b: Tensor, name: str = "not_equal") -> Tensor: + """Broadcasted element-wise comparison for (lhs != rhs). + + Parameters + ---------- + a : Tensor + The first input tensor. + + b : Tensor + The second input tensor. + + name : str + Name hint. + + Returns + ------- + result : Tensor + The computed result. + """ + return wrap_nested(_op.not_equal(a._expr, b._expr), name) + + +def cumsum( + data: Tensor, + axis: Optional[int] = None, + dtype: Optional[str] = None, + exclusive: Optional[bool] = None, + name: str = "cumsum", +) -> Tensor: + """Numpy style cumsum op. Return the cumulative inclusive sum of the elements along + a given axis. + + Parameters + ---------- + data : Tensor + The input data to the operator. + + axis : Optional[int] + Axis along which the cumulative sum is computed. The default (None) is to compute + the cumsum over the flattened array. + + dtype : Optional[str] + Type of the returned array and of the accumulator in which the elements are summed. + If dtype is not specified, it defaults to the dtype of data. + + exclusive : Optional[bool] + If true will return exclusive sum in which the first element is not + included. + + name : str + Name hint. + + Returns + ------- + result : Tensor + The result has the same size as data, and the same shape as data if axis is not None. + If axis is None, the result is a 1-d array. + + Examples + -------- + .. code-block:: python + + a = [[1, 2, 3], [4, 5, 6]] + + cumsum(a) # if axis is not provided, cumsum is done over the flattened input. + -> [ 1, 3, 6, 10, 15, 21] + + cumsum(a, dtype="float32") + -> [ 1., 3., 6., 10., 15., 21.] + + cumsum(a, axis=0) # sum over rows for each of the 3 columns + -> [[1, 2, 3], + [5, 7, 9]] + + cumsum(a, axis=1) + -> [[ 1, 3, 6], + [ 4, 9, 15]] + + a = [1, 0, 1, 0, 1, 1, 0] # a is a boolean array + cumsum(a, dtype=int32) # dtype should be provided to get the expected results + -> [1, 1, 2, 2, 3, 4, 4] + """ + return wrap_nested(_op.cumsum(data._expr, axis, dtype, exclusive), name) + + +def multinomial_from_uniform(prob: Tensor, uniform_sample: Tensor): + """Returns a tensor where each row contains the index sampled from the multinomial + probability distribution located in the corresponding row of tensor prob. + + Notes + ----- + For better cpu performance, use 'vm.builtin.multinomial_from_uniform'. + + Parameters + ---------- + prob : Tensor + The probability tensor. + + uniform_sample : Tensor + The uniform sample. + + Returns + ------- + result : Tensor + The computed result. + + Examples + -------- + .. code-block:: python + + prob = [[0.2, 0.3, 0.5], [0.3, 0.4, 0.3]] + usample = [[0.4], [0.9]] + + multinomial_from_uniform(a) + -> [[1], [2]] + """ + + batch = prob.shape[0] + cumsum_prob = cumsum(prob, axis=1, exclusive=False) + + @T.prim_func(private=True) + def _get_sample_index(A: T.handle, B: T.handle, C: T.handle): + batch, vocab_size = T.int64(), T.int64() + prob = T.match_buffer(A, (batch, vocab_size), "float32") + usample = T.match_buffer(B, (batch, 1), "float32") + output_index = T.match_buffer(C, (batch, 1), "int64") + + for ax0 in T.parallel(batch): + for ax1 in T.parallel(vocab_size): + with T.block("T_get_sample_index"): + v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) + T.writes(output_index[v_ax0, 0]) + if usample[v_ax0, T.int64(0)] < prob[v_ax0, v_ax1] or v_ax1 + 1 == vocab_size: + if v_ax1 == 0: + output_index[v_ax0, 0] = 0 + else: + if not (usample[v_ax0, T.int64(0)] < prob[v_ax0, v_ax1 - 1]): Review Comment: They are logically equivalent. Just compared the cuda code for the two, the codes are also equivalent. -- 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. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org