codeislife99 commented on a change in pull request #7149:
URL: https://github.com/apache/tvm/pull/7149#discussion_r551089181
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,59 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseSegmentSumAttrs);
+
+bool SparseSegmentSumRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [data, indices, segment_ids, result]
+ ICHECK_EQ(types.size(), 4) << "SparseSegmentSumRel expects 4 types but
provided " << types.size();
+ auto data = types[0].as<TensorTypeNode>();
+ auto indices = types[1].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ Array<PrimExpr> new_data_shape;
+ new_data_shape.push_back(tvm::max(indices->shape[0], param->num_segments));
+ for (int i = 1; i < static_cast<int>(data->shape.size()); ++i) {
+ new_data_shape.push_back(data->shape[i]);
+ }
+ std::vector<Type> fields;
+ fields.push_back(TensorType(new_data_shape, data->dtype));
+ fields.push_back(TensorType(Array<PrimExpr>{1}, tvm::DataType::Int(32)));
+ reporter->Assign(types[3], TupleType(Array<Type>(fields)));
+ return true;
+}
+
+Array<te::Tensor> SparseSegmentSumCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseSegmentSumCompute expects 3 input but
provided "
+ << inputs.size();
+ const auto* param = attrs.as<SparseSegmentSumAttrs>();
+ ICHECK(param != nullptr);
+ return {topi::SparseSegmentSum(inputs[0], inputs[1], inputs[2],
param->num_segments)};
+}
+
+Expr MakeSparseSegmentSum(Expr data, Expr indices, Expr segment_ids, int
num_segments) {
+ auto attrs = make_object<SparseSegmentSumAttrs>();
+ attrs->num_segments = std::move(num_segments);
+ static const Op& op = Op::Get("sparse_segment_sum");
+ return Call(op, {data, indices, segment_ids}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op._make.sparse_segment_sum").set_body_typed(MakeSparseSegmentSum);
+
+RELAY_REGISTER_OP("sparse_segment_sum")
+ .describe(R"code(Return sparse segment sum of the tensor given segments
+)code" TVM_ADD_FILELINE)
+ .set_num_inputs(3)
+ .set_attrs_type<SparseSegmentSumAttrs>()
+ .add_argument("data", "Tensor", "The first tensor")
+ .add_argument("indices", "Tensor", "The second tensor")
+ .add_argument("segment_ids", "Tensor", "The third tensor")
Review comment:
Done.
----------------------------------------------------------------
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]
