shingjan commented on code in PR #12485:
URL: https://github.com/apache/tvm/pull/12485#discussion_r959839869
##########
include/tvm/topi/transform.h:
##########
@@ -2016,6 +2016,70 @@ inline Tensor adv_index(const Tensor& data, const
Array<Tensor>& indices,
name, tag);
}
+/*!
+ * \brief Returns the sums, means or maxes of bags of embeddings
+ * \param input 1-d tensor of indics.
+ * \param weight the 2-d lookup table.
+ * \param offset the starting index position of each bag.
+ * \param mode the way (`sum`, `mean` or `max`) to reduce the bag.
+ * \param per_sample_weights the weights for every input.
+ * \param padding_idx the index of padding_idx is not contributed to the
result.
+ * \param include_last_offset if the last element of offset array is included.
+ * \param dtype the type of result.
+ * \param name output tensor name.
+ * \param tag output tensor tag.
+ * \return The embedded tensor.
+ */
+inline Tensor embedding_bag(const Tensor& input, const Tensor& weight, const
Tensor& offset,
+ int mode, const Tensor& per_sample_weights,
Integer padding_idx,
+ bool include_last_offset, DataType dtype,
+ const std::string name = "T_embedding_bag",
+ const std::string tag = kInjective) {
+ auto row = offset->shape[0];
+ auto column = weight->shape[1];
+
+ int N = GetConstInt(input->shape[0]);
+ if (include_last_offset) {
+ row = row - 1;
+ }
+
+ Array<PrimExpr> oshape{row, column};
+
+ auto func = [&](tvm::tir::Var i, tvm::tir::Var j) {
+ auto ret = make_zero(dtype);
+ auto count = make_zero(dtype); // count how many elements are used
+
+ auto st = offset(i); // start point
+ auto ed = tvm::tir::Select(row == i + 1, PrimExpr(N),
+ cast(DataType::Int(32), offset(i + 1))); //
end point
+
+ // can't find `fold` function, so use a stupid method to iterate from st
to ed
Review Comment:
lets make this comment more formal. Something like:
``` Use loop iteration here for the lack of fold in relay```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
