tkonolige commented on a change in pull request #6580:
URL: https://github.com/apache/incubator-tvm/pull/6580#discussion_r501889549
##########
File path: python/tvm/topi/nn/sparse.py
##########
@@ -207,3 +207,28 @@ def _csr_transpose_ir(data, indices, indptr, out_data,
out_indices, out_indptr):
last[0] = temp2[0]
return irb.get()
+
+
[email protected]_func
+def sparse_dense_alter_layout(_attrs, _inputs, _tinfos, _out_type):
Review comment:
Pylint complains these inputs are unused. Prefixing them with _ silences
the warning and indicates to other people that the argument is unused.
##########
File path: src/relay/op/nn/sparse.cc
##########
@@ -76,7 +76,36 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense")
});
RELAY_REGISTER_OP("nn.sparse_dense")
- .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T`
with X sparse.
+ .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T`
with W sparse.
+
+- **data**: `(x1, x2, ..., xn, input_dim)`
+- **weight**: `(units, input_dim)`
+- **out**: `(x1, x2, ..., xn, units)`.
+
+)code" TVM_ADD_FILELINE)
+ .set_attrs_type<SparseDenseAttrs>()
+ .set_num_inputs(4)
+ .add_argument("data", "nD Tensor", "Input data.")
+ .add_argument("weight_data", "1D Tensor", "Weight data matrix.")
+ .add_argument("weight_indices", "1D Tensor", "Weight indices matrix.")
+ .add_argument("weight_indptr", "1D Tensor", "Weight indptr matrix.")
+ .set_support_level(1)
+ .add_type_rel("SparseDense", SparseDenseRel);
+
+Expr MakeSparseDensePadded(Expr data, Expr weight_data, Expr weight_indices,
Expr weight_indptr) {
+ auto attrs = make_object<SparseDenseAttrs>();
+ static const Op& op = Op::Get("nn.sparse_dense_padded");
+ return Call(op, {data, weight_data, weight_indices, weight_indptr},
Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense_padded")
+ .set_body([](const TVMArgs& args, TVMRetValue* rv) {
+ runtime::detail::unpack_call<Expr, 4>(MakeSparseDensePadded, args, rv);
+ });
+
+RELAY_REGISTER_OP("nn.sparse_dense_padded")
Review comment:
This op shouldn't be used directly by a user unless they are sure that
the input matrix is in the correct format. I've documented this.
##########
File path: src/relay/op/nn/sparse.cc
##########
@@ -76,7 +76,36 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense")
});
RELAY_REGISTER_OP("nn.sparse_dense")
- .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T`
with X sparse.
+ .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T`
with W sparse.
+
+- **data**: `(x1, x2, ..., xn, input_dim)`
+- **weight**: `(units, input_dim)`
+- **out**: `(x1, x2, ..., xn, units)`.
+
+)code" TVM_ADD_FILELINE)
+ .set_attrs_type<SparseDenseAttrs>()
+ .set_num_inputs(4)
+ .add_argument("data", "nD Tensor", "Input data.")
+ .add_argument("weight_data", "1D Tensor", "Weight data matrix.")
+ .add_argument("weight_indices", "1D Tensor", "Weight indices matrix.")
+ .add_argument("weight_indptr", "1D Tensor", "Weight indptr matrix.")
+ .set_support_level(1)
+ .add_type_rel("SparseDense", SparseDenseRel);
+
+Expr MakeSparseDensePadded(Expr data, Expr weight_data, Expr weight_indices,
Expr weight_indptr) {
+ auto attrs = make_object<SparseDenseAttrs>();
+ static const Op& op = Op::Get("nn.sparse_dense_padded");
+ return Call(op, {data, weight_data, weight_indices, weight_indptr},
Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense_padded")
+ .set_body([](const TVMArgs& args, TVMRetValue* rv) {
+ runtime::detail::unpack_call<Expr, 4>(MakeSparseDensePadded, args, rv);
+ });
+
+RELAY_REGISTER_OP("nn.sparse_dense_padded")
Review comment:
How can I not expose it? Do I just change the name to
"internal.nn.sparse_dense_padded"?
##########
File path: src/relay/op/nn/sparse.cc
##########
@@ -76,7 +76,36 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense")
});
RELAY_REGISTER_OP("nn.sparse_dense")
- .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T`
with X sparse.
+ .describe(R"code(Applies a sparse linear transformation: :math:`Y = XW^T`
with W sparse.
+
+- **data**: `(x1, x2, ..., xn, input_dim)`
+- **weight**: `(units, input_dim)`
+- **out**: `(x1, x2, ..., xn, units)`.
+
+)code" TVM_ADD_FILELINE)
+ .set_attrs_type<SparseDenseAttrs>()
+ .set_num_inputs(4)
+ .add_argument("data", "nD Tensor", "Input data.")
+ .add_argument("weight_data", "1D Tensor", "Weight data matrix.")
+ .add_argument("weight_indices", "1D Tensor", "Weight indices matrix.")
+ .add_argument("weight_indptr", "1D Tensor", "Weight indptr matrix.")
+ .set_support_level(1)
+ .add_type_rel("SparseDense", SparseDenseRel);
+
+Expr MakeSparseDensePadded(Expr data, Expr weight_data, Expr weight_indices,
Expr weight_indptr) {
+ auto attrs = make_object<SparseDenseAttrs>();
+ static const Op& op = Op::Get("nn.sparse_dense_padded");
+ return Call(op, {data, weight_data, weight_indices, weight_indptr},
Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op.nn._make.sparse_dense_padded")
+ .set_body([](const TVMArgs& args, TVMRetValue* rv) {
+ runtime::detail::unpack_call<Expr, 4>(MakeSparseDensePadded, args, rv);
+ });
+
+RELAY_REGISTER_OP("nn.sparse_dense_padded")
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]
