tkonolige commented on a change in pull request #7126:
URL: https://github.com/apache/tvm/pull/7126#discussion_r547476827
##########
File path: python/tvm/relay/op/transform.py
##########
@@ -1320,3 +1320,83 @@ def adv_index(inputs):
Output tensor.
"""
return _make.adv_index(Tuple(inputs))
+
+
+def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape,
default_value):
+ """
+ Fill first column of the empty rows with default values for a sparse array.
+ It returns a TupleWrapper with four outputs
+
+ Parameters
+ ----------
+ sparse_indices : relay.Expr
+ A 2-D tensor[N, n_dim] of integers containing location of sparse
values, where N is the
+ number of sparse values and n_dim is the number of dimensions of the
dense_shape
+
+ sparse_values : relay.Expr
+ A 1-D tensor[N] containing the sparse values for the sparse indices.
+
+ dense_shape : relay.Expr
+ A list of integers. Shape of the dense output tensor.
+
+ default_value : relay.Expr
+ A 0-D tensor containing the default value for the remaining locations.
+ Defaults to 0.
+
+ Returns
+ -------
+ new_sparse_indices : relay.Expr
+ A 2-D tensor[N + dense_shape[0], n_dim] of integers containing
location of new sparse
+ indices where N is the number of sparse values. It is filled with -1
at irrelevant indices
+ which will be sliced in a future op discarding non-useful elements.
This is done since the
+ real rows of new_sparse_indices depends on the input.
+
+ empty_row_indicator : relay.Expr
+ A 1-D Boolean tensor[dense_shape[0]] indicating whether the particular
row is empty
+
+ new_sparse_values : relay.Expr
+ A 1-D tensor[dense_shape[0]] containing the sparse values for the
sparse indices. It is
+ filled with -1 at to_be_discarded indices
Review comment:
Could you update the wording here? `to_be_discarded` refers to nothing.
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,65 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseFillEmptyRowsAttrs);
+
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [ sparse_indices, sparse_values, default_values, result]
+ ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments
but provided "
+ << types.size();
+ std::vector<Type> fields;
+ auto sparse_indices = types[0].as<TensorTypeNode>();
+ auto default_value = types[2].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+ ICHECK(param != nullptr);
+
+ Array<IndexExpr> sp_ordered_output_shape;
+ sp_ordered_output_shape.push_back(param->dense_shape[0] +
sparse_indices->shape[0]);
+ if (sparse_indices->shape.size() > 1) {
+ sp_ordered_output_shape.push_back(sparse_indices->shape[1]);
+ }
+ fields.push_back(TensorType(sp_ordered_output_shape, sparse_indices->dtype));
+ fields.push_back(TensorType(Array<PrimExpr>{param->dense_shape[0]},
tvm::DataType::Bool()));
+ fields.push_back(TensorType(Array<PrimExpr>{sp_ordered_output_shape[0]},
default_value->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> SparseFillEmptyRowsCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseFillEmptyRowsCompute expects 3
arguments but provided "
+ << inputs.size();
+ const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+ ICHECK(param != nullptr);
Review comment:
```suggestion
ICHECK_NOTNULL(param);
```
##########
File path: include/tvm/topi/transform.h
##########
@@ -1386,6 +1386,96 @@ inline Array<Tensor> meshgrid(const Array<Tensor>&
inputs, const std::string& in
return result;
}
+/*!
+ * \brief Fill Empty rows of a sparse tensor with default value
Review comment:
```suggestion
* \brief Fill empty rows of a sparse tensor with default values
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,65 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseFillEmptyRowsAttrs);
+
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [ sparse_indices, sparse_values, default_values, result]
+ ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments
but provided "
+ << types.size();
+ std::vector<Type> fields;
+ auto sparse_indices = types[0].as<TensorTypeNode>();
+ auto default_value = types[2].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+ ICHECK(param != nullptr);
Review comment:
```suggestion
ICHECK_NOTNULL(param);
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,65 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseFillEmptyRowsAttrs);
+
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [ sparse_indices, sparse_values, default_values, result]
+ ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments
but provided "
+ << types.size();
Review comment:
```suggestion
ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments
but " << types.size() << " were provided.";
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,65 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseFillEmptyRowsAttrs);
+
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [ sparse_indices, sparse_values, default_values, result]
+ ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments
but provided "
+ << types.size();
+ std::vector<Type> fields;
+ auto sparse_indices = types[0].as<TensorTypeNode>();
+ auto default_value = types[2].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+ ICHECK(param != nullptr);
+
+ Array<IndexExpr> sp_ordered_output_shape;
+ sp_ordered_output_shape.push_back(param->dense_shape[0] +
sparse_indices->shape[0]);
+ if (sparse_indices->shape.size() > 1) {
+ sp_ordered_output_shape.push_back(sparse_indices->shape[1]);
+ }
+ fields.push_back(TensorType(sp_ordered_output_shape, sparse_indices->dtype));
+ fields.push_back(TensorType(Array<PrimExpr>{param->dense_shape[0]},
tvm::DataType::Bool()));
+ fields.push_back(TensorType(Array<PrimExpr>{sp_ordered_output_shape[0]},
default_value->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> SparseFillEmptyRowsCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseFillEmptyRowsCompute expects 3
arguments but provided "
+ << inputs.size();
Review comment:
```suggestion
ICHECK_EQ(inputs.size(), 3) << "SparseFillEmptyRowsCompute expects 3
arguments but " << inputs.size() << " were provided.";
```
##########
File path: src/relay/op/tensor/transform.cc
##########
@@ -1553,6 +1553,65 @@ RELAY_REGISTER_OP("meshgrid")
.set_attr<FTVMCompute>("FTVMCompute", MeshgridCompute)
.set_attr<TOpPattern>("TOpPattern", kInjective);
+TVM_REGISTER_NODE_TYPE(SparseFillEmptyRowsAttrs);
+
+bool SparseFillEmptyRowsRel(const Array<Type>& types, int num_inputs, const
Attrs& attrs,
+ const TypeReporter& reporter) {
+ // types: [ sparse_indices, sparse_values, default_values, result]
+ ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments
but provided "
+ << types.size();
+ std::vector<Type> fields;
+ auto sparse_indices = types[0].as<TensorTypeNode>();
+ auto default_value = types[2].as<TensorTypeNode>();
+ const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+ ICHECK(param != nullptr);
+
+ Array<IndexExpr> sp_ordered_output_shape;
+ sp_ordered_output_shape.push_back(param->dense_shape[0] +
sparse_indices->shape[0]);
+ if (sparse_indices->shape.size() > 1) {
+ sp_ordered_output_shape.push_back(sparse_indices->shape[1]);
+ }
+ fields.push_back(TensorType(sp_ordered_output_shape, sparse_indices->dtype));
+ fields.push_back(TensorType(Array<PrimExpr>{param->dense_shape[0]},
tvm::DataType::Bool()));
+ fields.push_back(TensorType(Array<PrimExpr>{sp_ordered_output_shape[0]},
default_value->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> SparseFillEmptyRowsCompute(const Attrs& attrs, const
Array<te::Tensor>& inputs,
+ const Type& out_type) {
+ ICHECK_EQ(inputs.size(), 3) << "SparseFillEmptyRowsCompute expects 3
arguments but provided "
+ << inputs.size();
+ const auto* param = attrs.as<SparseFillEmptyRowsAttrs>();
+ ICHECK(param != nullptr);
+ return {topi::SparseFillEmptyRows(inputs[0], inputs[1], inputs[2],
param->dense_shape)};
+}
+
+Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr
default_value,
+ Array<Integer> dense_shape) {
+ auto attrs = make_object<SparseFillEmptyRowsAttrs>();
+ attrs->dense_shape = std::move(dense_shape);
+ static const Op& op = Op::Get("sparse_fill_empty_rows");
+ return Call(op, {sparse_indices, sparse_values, default_value},
Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op._make.sparse_fill_empty_rows")
+ .set_body_typed(MakeSparseFillEmptyRows);
+
+RELAY_REGISTER_OP("sparse_fill_empty_rows")
+ .describe(R"code(Return representation of a sparse tensor with empty rows
filled with default
+ value.)code" TVM_ADD_FILELINE)
+ .set_num_inputs(3)
+ .set_attrs_type<SparseFillEmptyRowsAttrs>()
+ .add_argument("sparse_indices", "Tensor", "The first tensor")
Review comment:
Please put more useful descriptions on these arguments.
----------------------------------------------------------------
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]
