From a14842cc8c602e17c61a482eb1c3031eaee2b53b Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 02:46:10 +0000 Subject: [PATCH 01/32] Initial Commit --- python/tvm/relay/op/_transform.py | 66 +++++++++++ python/tvm/relay/op/strategy/generic.py | 29 +++++ python/tvm/relay/op/transform.py | 6 + python/tvm/topi/__init__.py | 1 + python/tvm/topi/generic/search.py | 4 + python/tvm/topi/sparse_fill_empty_rows.py | 105 +++++++++++++++++ python/tvm/topi/transform.py | 4 + src/relay/op/tensor/transform.cc | 47 ++++++++ .../relay/dyn/test_dynamic_op_level3.py | 108 +++++++++++++++++- 9 files changed, 367 insertions(+), 3 deletions(-) create mode 100644 python/tvm/topi/sparse_fill_empty_rows.py diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index ba2416ff8950..e737361567cd 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -94,6 +94,24 @@ def compute_scatter(attrs, inputs, output_type): _reg.register_strategy("scatter", strategy.scatter_strategy) +# sparse_fill_empty_rows +@_reg.register_compute("sparse_fill_empty_rows") +def compute_sparse_fill_empty_rows(attrs, inputs, output_type): + """Compute definition of sparse_fill_empty_rows""" + + return topi.sparse_fill_empty_rows( + inputs[0], + inputs[1], + inputs[2], + inputs[3], + output_type.fields[0].shape, + output_type.fields[1].shape, + output_type.fields[2].shape, + ) + + +_reg.register_strategy("sparse_fill_empty_rows", strategy.sparse_fill_empty_rows_strategy) + # scatter_add @_reg.register_compute("scatter_add") def compute_scatter_add(attrs, inputs, output_type): @@ -445,6 +463,54 @@ def argwhere_shape_func(attrs, inputs, out_ndims): _reg.register_shape_func("scatter_add", False, elemwise_shape_func) +# @script +# def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): +# out = output_tensor((2,), "int64") +# out2 = output_tensor((1,), "int64") +# num_rows = int64(dense_shape[0]) +# count = int64(sparse_indices.shape[0]) +# for i in const_range(1, sparse_indices.shape[0]): +# index = sparse_indices[i, 0] +# prev_index = sparse_indices[i - 1, 0] + 1 + +# if int64(index) > int64(prev_index): +# count += int64(index) - int64(prev_index) + +# count += int64(num_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]) +# out[0] = int64(count) +# out[1] = int64(2) +# out2[0] = int64(count) +# return (out, out2) + + +@script +def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): + new_sparse_indices_shape = output_tensor((2,), "int64") + new_sparse_values_shape = output_tensor((1,), "int64") + empty_row_indicator_shape = output_tensor((1,), "int64") + num_rows = int64(dense_shape[0]) + count = int64(sparse_indices.shape[0]) + for i in const_range(1, sparse_indices.shape[0]): + index = int64(sparse_indices[i, 0]) + prev_index = int64(sparse_indices[i - 1, 0] + 1) + + if index > prev_index: + count += index - prev_index + + count += int64(sparse_indices[0, 0]) + count += int64(num_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]) + new_sparse_indices_shape[0] = int64(count) + new_sparse_indices_shape[1] = int64(sparse_indices.shape[1]) + new_sparse_values_shape[0] = int64(count) + empty_row_indicator_shape[0] = int64(dense_shape[0]) + return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape) + + +@_reg.register_shape_func("sparse_fill_empty_rows", True) +def sparse_fill_empty_rows_func(attrs, inputs, _): + return _sparse_fill_empty_rows_shape_func(inputs[0], inputs[2]) + + @script def _layout_transform_shape_func( data_shape, out_layout_len, dst_equal_list, dst_mul_list, dst_div_list, dst_mix_list diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py index af1d2552fab7..d4b1d15a5318 100644 --- a/python/tvm/relay/op/strategy/generic.py +++ b/python/tvm/relay/op/strategy/generic.py @@ -1055,6 +1055,35 @@ def roi_align_strategy(attrs, inputs, out_type, target): return strategy +# sparse_fill_empty_rows +@override_native_generic_func("sparse_fill_empty_rows_strategy") +def sparse_fill_empty_rows_strategy(attrs, outs, out_type, target): + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_sparse_fill_empty_rows(topi.sparse_fill_empty_rows), + wrap_topi_schedule(topi.generic.schedule_sparse_fill_empty_rows), + name="sparse_fill_empty_rows.generic", + ) + return strategy + + +def wrap_compute_sparse_fill_empty_rows(topi_compute): + """Wrap sparse_fill_empty_rows compute""" + + def _compute_sparse_fill_empty_rows(attrs, inputs, output_type): + return topi_compute( + inputs[0], + inputs[1], + inputs[2], + inputs[3], + output_type.fields[0].shape, + output_type.fields[1].shape, + output_type.fields[2].shape, + ) + + return _compute_sparse_fill_empty_rows + + # roi_pool @generic_func def schedule_roi_pool(attrs, outs, target): diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index e9d081eb5fb6..184bc7156712 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1322,6 +1322,12 @@ def adv_index(inputs): return _make.adv_index(Tuple(inputs)) +def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): + return TupleWrapper( + _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 2 + ) + + def cumsum(data, axis=None, dtype=None, exclusive=None): """Numpy style cumsum op. Return the cumulative inclusive sum of the elements along a given axis. diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py index 6836f04b5ada..2b17162048e0 100644 --- a/python/tvm/topi/__init__.py +++ b/python/tvm/topi/__init__.py @@ -38,6 +38,7 @@ from .broadcast import * from .sort import * from .scatter import * +from .sparse_fill_empty_rows import * from .scatter_add import * from .argwhere import * from .cumsum import * diff --git a/python/tvm/topi/generic/search.py b/python/tvm/topi/generic/search.py index b3c8772046fd..a1767129bc24 100644 --- a/python/tvm/topi/generic/search.py +++ b/python/tvm/topi/generic/search.py @@ -66,3 +66,7 @@ def schedule_scatter_add(outs): The computation schedule for the op. """ return _default_schedule(outs, False) + + +def schedule_sparse_fill_empty_rows(outs): + return _default_schedule(outs, False) \ No newline at end of file diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py new file mode 100644 index 000000000000..15955b1244f9 --- /dev/null +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -0,0 +1,105 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHnew_sparse_indices WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks +"""SparseFillEmptyRows operator""" +from ..tir import decl_buffer, ir_builder, Cast, AssertStmt, StringImm, Evaluate +from ..te import extern, hybrid + + +@hybrid.script +def _sparse_fill_empty_rows( + sparse_indices, + sparse_values, + dense_shape, + default_value, + new_sparse_indices_shape, + new_sparse_values_shape, + empty_row_indicator_shape, +): + default_value_ = int64(default_value[0]) + new_sparse_indices = output_tensor(new_sparse_indices_shape, "int64") + new_sparse_values = output_tensor(new_sparse_values_shape, "int64") + empty_row_indicator = output_tensor(empty_row_indicator_shape, "int64") + idx = 0 + + for i in range(0, int64(sparse_indices[0, 0])): + new_sparse_indices[idx, 0] = int64(i) + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(0) + + new_sparse_values[idx] = default_value_ + empty_row_indicator[i] = int64(1) + idx += 1 + + for i in const_range(0, sparse_indices.shape[0]): + index = int64(sparse_indices[i, 0]) + if i == 0: + new_sparse_indices[idx, 0] = index + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) + new_sparse_values[idx] = int64(sparse_values[i]) + empty_row_indicator[index] = int64(0) + idx += 1 + else: + prev_index = int64(sparse_indices[i - 1, 0] + 1) + for j in range(prev_index, index): + new_sparse_indices[idx, 0] = int64(j) + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(0) + empty_row_indicator[prev_index] = int64(1) + new_sparse_values[idx] = default_value_ + idx += 1 + + new_sparse_indices[idx, 0] = index + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) + new_sparse_values[idx] = int64(sparse_values[i]) + empty_row_indicator[index] = int64(0) + idx += 1 + + for i in range( + int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1), int64(dense_shape[0]) + ): + + new_sparse_indices[idx, 0] = int64(i) + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(0) + empty_row_indicator[i] = int64(1) + new_sparse_values[idx] = default_value_ + idx += 1 + + return (new_sparse_indices, new_sparse_values, empty_row_indicator) + + +def sparse_fill_empty_rows( + sparse_indices, + sparse_values, + dense_shape, + default_value, + new_sparse_indices_shape, + new_sparse_values_shape, + empty_row_indicator_shape, +): + return _sparse_fill_empty_rows( + sparse_indices, + sparse_values, + dense_shape, + default_value, + new_sparse_indices_shape, + new_sparse_values_shape, + empty_row_indicator_shape, + ) diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index 6ddbc73e4666..308004e2a798 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -931,3 +931,7 @@ def adv_index(data, indices): Output tensor """ return cpp.adv_index(data, indices) + + +def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): + return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) \ No newline at end of file diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 5e39b409615d..1c332deac73d 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1584,6 +1584,53 @@ RELAY_REGISTER_OP("repeat") .set_attr("FTVMCompute", RepeatCompute) .set_attr("TOpPattern", kBroadcast); +bool SparseFillEmptyRowsRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + // types: [sparse_indices, sparse_values, dense_shape, default_value, result] + ICHECK_EQ(types.size(), 5); + std::vector fields; + auto sparse_indices = types[0].as(); + auto ndims = sparse_indices->shape[1]; + // auto dense_shape = types[2].as(); + fields.push_back(TensorType(Array{Any(), ndims}, tvm::DataType::Int(64))); + // fields.push_back(TensorType(Array{Any()}, tvm::DataType::Int(64))); + + fields.push_back(TensorType(Array{Any()}, tvm::DataType::Int(64))); + fields.push_back(TensorType(Array{Any()}, tvm::DataType::Int(64))); + + // auto sample_input = types[0].as(); + reporter->Assign(types[types.size() - 1], TupleType(Array(fields))); + + // reporter->Assign(types[types.size()-1], TensorType(Array{Any()}, tvm::DataType::Int(64))); + return true; +} + +Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr dense_shape, Expr default_value) { + static const Op& op = Op::Get("sparse_fill_empty_rows"); + return Call(op, {sparse_indices, sparse_values, dense_shape, default_value}, 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(4) + .add_argument("sparse_indices", "Tensor", + "A 2-D int64 tensor of shape [N, ndims], which specifies the indices of the" + "elements in the sparse tensor that contain nonzero values") + .add_argument("sparse_values", "Tensor", + "A 1-D tensor[N] which supplies the values for each element in indices") + .add_argument("dense_shape", "Tensor", + "A 1-D int64 tensor of shape [ndims], which specifies the dense_shape of the" + "sparse tensor. Takes a list indicating the number of elements in each dimension") + .add_argument("default_value", "Tensor", + "The value to fill for empty rows, with the same type as sparse_values") + .add_type_rel("sparse_fill_empty_rows", SparseFillEmptyRowsRel) + .set_support_level(3) + .set_attr("TOpPattern", kOpaque); + // meshgrid operator TVM_REGISTER_NODE_TYPE(MeshgridAttrs); diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index dd73b9a96a52..fa993e61b577 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -29,11 +29,14 @@ def verify_func(func, data, ref_res): assert isinstance(data, list) for target, ctx in tvm.testing.enabled_targets(): - for kind in ["vm", "debug"]: + for kind in ["vm"]: mod = tvm.ir.IRModule.from_expr(func) intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target) op_res = intrp.evaluate()(*data) - tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5) + for op_result, ref_result in zip(op_res, ref_res): + # print(op_result.asnumpy(), ref_result) + print(op_result.asnumpy()) + # tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5) relay.backend.compile_engine.get().clear() @@ -202,5 +205,104 @@ def verify_sparse_to_dense(sparse_indices, sparse_values, default_value, output_ verify_sparse_to_dense(1, 3, None, [5], [0, 3, 0, 0, 0]) # default value not specified +def test_sparse_fill_empty_rows(): + def ref_sparse_fill_empty_rows( + sparse_indices: np.ndarray, + sparse_values: np.ndarray, + dense_shape: np.ndarray, + default_value: np.ndarray, + ) -> None: + """ + This function calculates the expected output of sparse_fill_empty_rows operator given the + inputs. + """ + new_sparse_indices = -1 * np.ones( + (sparse_indices.shape[0] + dense_shape[0], sparse_indices.shape[1]) + ) + empty_row_indicator = np.ones(dense_shape[0], dtype=bool) + new_sparse_values = -1 * np.ones(sparse_indices.shape[0] + dense_shape[0]) + + for i in range(sparse_indices.shape[0]): + empty_row_indicator[sparse_indices[i][0]] = False + + new_sparse_indices[: sparse_indices.shape[0]][:] = sparse_indices[:] + new_sparse_values[: sparse_values.shape[0]] = sparse_values[:] + new_sparse_indices_index = sparse_indices.shape[0] + + for empty_row_index, element in enumerate(empty_row_indicator): + if element: + new_sparse_indices[new_sparse_indices_index, 0] = empty_row_index + new_sparse_indices[new_sparse_indices_index, 1:] = 0 + new_sparse_values[new_sparse_indices_index] = default_value[0] + new_sparse_indices_index += 1 + + return new_sparse_indices, new_sparse_values, empty_row_indicator + + def verify_sparse_fill_empty_rows( + sparse_indices_np: np.ndarray, + sparse_values_np: np.ndarray, + dense_shape_np: np.ndarray, + default_value_np: np.ndarray, + ) -> None: + """ + This function verifies the relay output of sparse_fill_empty_rows with its expected output. + """ + sparse_indices = relay.var( + "sparse_indices", + relay.TensorType(sparse_indices_np.shape, str(sparse_indices_np.dtype)), + ) + sparse_values = relay.var( + "sparse_values", + relay.TensorType(sparse_values_np.shape, str(sparse_values_np.dtype)), + ) + dense_shape = relay.var( + "dense_shape", + relay.TensorType(dense_shape_np.shape, str(dense_shape_np.dtype)), + ) + default_value = relay.var( + "default_value", + relay.TensorType(default_value_np.shape, str(default_value_np.dtype)), + ) + z = relay.sparse_fill_empty_rows( + sparse_indices, sparse_values, dense_shape, default_value + ).astuple() + func = relay.Function([sparse_indices, sparse_values, dense_shape, default_value], z) + ref_res = ref_sparse_fill_empty_rows( + sparse_indices_np, + sparse_values_np, + dense_shape_np, + default_value_np, + ) + # ref_res = [] + verify_func( + func, [sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np], ref_res + ) + + sparse_indices = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int32) + sparse_values = np.array([1, 2, 3, 4], dtype=np.int32) + dense_shape = np.array([7, 6], dtype=np.int32) + default_value = np.array([5], dtype=sparse_values.dtype) + verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) + + sparse_indices = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int32) + sparse_values = np.array([1, 2, 3, 4], dtype=np.int32) + dense_shape = np.array([7, 7, 7], dtype=np.int32) + default_value = np.array([10], dtype=np.int32) + verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) + + sparse_indices = np.array([[1], [2]], dtype=np.int32) + sparse_values = np.array([7, 8], dtype=np.int32) + dense_shape = np.array([5], dtype=np.int32) + default_value = np.array([4], dtype=np.int32) + verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) + + # sparse_indices = np.ones((0, 1), dtype=np.int32) + # sparse_values = np.array([], dtype=np.int32) + # dense_shape = np.array([5], dtype=np.int32) + # default_value = np.array([4], dtype=np.int32) + # verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) + + if __name__ == "__main__": - pytest.main([__file__]) + test_sparse_fill_empty_rows() + # pytest.main([__file__]) From cf4f0bdae27b58193e2fd78d425d7ec2a221f7a0 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 02:48:52 +0000 Subject: [PATCH 02/32] Fix formats --- python/tvm/relay/op/_transform.py | 20 -------------------- src/relay/op/tensor/transform.cc | 8 +++++--- 2 files changed, 5 insertions(+), 23 deletions(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index e737361567cd..760b88891d50 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -463,26 +463,6 @@ def argwhere_shape_func(attrs, inputs, out_ndims): _reg.register_shape_func("scatter_add", False, elemwise_shape_func) -# @script -# def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): -# out = output_tensor((2,), "int64") -# out2 = output_tensor((1,), "int64") -# num_rows = int64(dense_shape[0]) -# count = int64(sparse_indices.shape[0]) -# for i in const_range(1, sparse_indices.shape[0]): -# index = sparse_indices[i, 0] -# prev_index = sparse_indices[i - 1, 0] + 1 - -# if int64(index) > int64(prev_index): -# count += int64(index) - int64(prev_index) - -# count += int64(num_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]) -# out[0] = int64(count) -# out[1] = int64(2) -# out2[0] = int64(count) -# return (out, out2) - - @script def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): new_sparse_indices_shape = output_tensor((2,), "int64") diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 1c332deac73d..62dcd8dd164b 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1601,11 +1601,13 @@ bool SparseFillEmptyRowsRel(const Array& types, int num_inputs, const Attr // auto sample_input = types[0].as(); reporter->Assign(types[types.size() - 1], TupleType(Array(fields))); - // reporter->Assign(types[types.size()-1], TensorType(Array{Any()}, tvm::DataType::Int(64))); + // reporter->Assign(types[types.size()-1], TensorType(Array{Any()}, + // tvm::DataType::Int(64))); return true; } -Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr dense_shape, Expr default_value) { +Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr dense_shape, + Expr default_value) { static const Op& op = Op::Get("sparse_fill_empty_rows"); return Call(op, {sparse_indices, sparse_values, dense_shape, default_value}, Attrs(), {}); } @@ -1625,7 +1627,7 @@ RELAY_REGISTER_OP("sparse_fill_empty_rows") .add_argument("dense_shape", "Tensor", "A 1-D int64 tensor of shape [ndims], which specifies the dense_shape of the" "sparse tensor. Takes a list indicating the number of elements in each dimension") - .add_argument("default_value", "Tensor", + .add_argument("default_value", "Tensor", "The value to fill for empty rows, with the same type as sparse_values") .add_type_rel("sparse_fill_empty_rows", SparseFillEmptyRowsRel) .set_support_level(3) From e50c7f357c4c2a0f7d87511d09576b9c2a6d7528 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 02:53:32 +0000 Subject: [PATCH 03/32] Remove comments --- src/relay/op/tensor/transform.cc | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 62dcd8dd164b..02e4a5e1e113 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1591,18 +1591,10 @@ bool SparseFillEmptyRowsRel(const Array& types, int num_inputs, const Attr std::vector fields; auto sparse_indices = types[0].as(); auto ndims = sparse_indices->shape[1]; - // auto dense_shape = types[2].as(); fields.push_back(TensorType(Array{Any(), ndims}, tvm::DataType::Int(64))); - // fields.push_back(TensorType(Array{Any()}, tvm::DataType::Int(64))); - fields.push_back(TensorType(Array{Any()}, tvm::DataType::Int(64))); fields.push_back(TensorType(Array{Any()}, tvm::DataType::Int(64))); - - // auto sample_input = types[0].as(); reporter->Assign(types[types.size() - 1], TupleType(Array(fields))); - - // reporter->Assign(types[types.size()-1], TensorType(Array{Any()}, - // tvm::DataType::Int(64))); return true; } From 43fc9ac0f040ce561427fffb2dba69cf8e65fe93 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 02:57:18 +0000 Subject: [PATCH 04/32] Black --- python/tvm/topi/generic/search.py | 2 +- python/tvm/topi/transform.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/topi/generic/search.py b/python/tvm/topi/generic/search.py index a1767129bc24..5924d35def73 100644 --- a/python/tvm/topi/generic/search.py +++ b/python/tvm/topi/generic/search.py @@ -69,4 +69,4 @@ def schedule_scatter_add(outs): def schedule_sparse_fill_empty_rows(outs): - return _default_schedule(outs, False) \ No newline at end of file + return _default_schedule(outs, False) diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index 308004e2a798..0ac9bf9f9715 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -934,4 +934,4 @@ def adv_index(data, indices): def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): - return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) \ No newline at end of file + return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) From 09a8c5b37bd13cd1d4f2b4abd0cd69a38ce36020 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 03:13:24 +0000 Subject: [PATCH 05/32] THreeops --- python/tvm/relay/op/transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 184bc7156712..20e1d04f99d7 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1324,7 +1324,7 @@ def adv_index(inputs): def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): return TupleWrapper( - _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 2 + _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 3 ) From 6dba7f03ab674dae23cd57b83ccff480114eaa2e Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 03:44:36 +0000 Subject: [PATCH 06/32] Add Frontend Code --- python/tvm/relay/frontend/tensorflow.py | 18 ++++++ .../frontend/tensorflow/test_forward.py | 63 ++++++++++++++++++- 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index b34e6c723645..fd6505d3abdf 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -998,6 +998,23 @@ def _impl(inputs, attr, params, mod): return _impl +def _sparse_fill_empty_rows(): + def _impl(inputs, attr, params, mod): + assert len(inputs) == 4, "There should be 4 input tensors" + new_sparse_indices, new_sparse_values, empty_row_indicator = _op.sparse_fill_empty_rows( + inputs[0], inputs[1], inputs[2], inputs[3] + ) + + return _expr.TupleWrapper( + _expr.Tuple( + [new_sparse_indices, new_sparse_values, _op.cast(empty_row_indicator, dtype="bool")] + ), + 3, + ) + + return _impl + + def _identity(): def _impl(inputs, attr, params, mod): return inputs[0] @@ -2447,6 +2464,7 @@ def _impl(inputs, attr, params, mod): "SpaceToDepth": _space_to_depth(), "SparseToDense": _sparse_to_dense(), "SparseTensorDenseMatMul": _sparse_tensor_dense_matmul(), + "SparseFillEmptyRows": _sparse_fill_empty_rows(), "Split": _split(False), "SplitV": _split(True), "Sqrt": AttrCvt("sqrt"), diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 34ee0f3528ae..70ecf6f85ec5 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1812,6 +1812,66 @@ def test_forward_sparse_dense_matmul(): ) +####################################################################### +# SparseFillEmptyRows +# ------------ + + +def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_value): + with tf.Graph().as_default(): + indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices") + values = tf.placeholder(shape=values_np.shape, dtype=values_np.dtype, name="values") + dense_shape = tf.placeholder( + shape=dense_shape_np.shape, dtype=dense_shape_np.dtype, name="dense_shape" + ) + sp_input = tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=dense_shape) + _ = tf.sparse.fill_empty_rows(sp_input, default_value, name="sparse_fill_empty_rows") + compare_tf_with_tvm( + [indices_np, values_np, dense_shape_np], + [indices.name, values.name, dense_shape.name], + [ + "sparse_fill_empty_rows/SparseFillEmptyRows:0", + "sparse_fill_empty_rows/SparseFillEmptyRows:1", + "sparse_fill_empty_rows/SparseFillEmptyRows:2", + ], + mode="vm", + ) + + +def test_forward_sparse_fill_empty_rows(): + """ sparse_fill_empty_rows op test""" + ################################################################### + # + # In order to create a SparseTensor, it requires 3 input as below: + # SparseTensor(indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4]) + # + # Above Sparse can be represented in Dense as below : + # [[1, 0, 0, 0] + # [0, 0, 2, 0] + # [0, 0, 0, 0]] + # + # ------------------------------------------------------------------ + sparse_indices_np = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64) + sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int64) + dense_shape_np = np.array([5, 6], dtype=np.int64) + default_value = 10 + _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, dense_shape_np, default_value) + + sparse_indices_np = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64) + sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int64) + dense_shape_np = np.array([7, 7, 7], dtype=np.int64) + default_value_np = 5 + _test_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np + ) + + sparse_indices_np = np.array([[1], [2]], dtype=np.int64) + sparse_values_np = np.array([7, 8], dtype=np.int64) + dense_shape_np = np.array([5], dtype=np.int64) + default_value_np = 4 + _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, dense_shape_np, default_value) + + ####################################################################### # StridedSlice # ------------ @@ -4687,4 +4747,5 @@ def lstm_cell(): if __name__ == "__main__": - pytest.main([__file__]) + test_forward_sparse_fill_empty_rows() + # pytest.main([__file__]) From d593612241e302a1b880cf6d5f420c468a41cd13 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 03:48:02 +0000 Subject: [PATCH 07/32] Add Default Value to feed dict --- tests/python/frontend/tensorflow/test_forward.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 70ecf6f85ec5..0fee664b7dfd 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1817,18 +1817,19 @@ def test_forward_sparse_dense_matmul(): # ------------ -def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_value): +def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_value_int): with tf.Graph().as_default(): indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices") values = tf.placeholder(shape=values_np.shape, dtype=values_np.dtype, name="values") dense_shape = tf.placeholder( shape=dense_shape_np.shape, dtype=dense_shape_np.dtype, name="dense_shape" ) + default_value = tf.placeholder(shape=(), dtype = values_np.dtype, name="default_value") sp_input = tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=dense_shape) _ = tf.sparse.fill_empty_rows(sp_input, default_value, name="sparse_fill_empty_rows") compare_tf_with_tvm( - [indices_np, values_np, dense_shape_np], - [indices.name, values.name, dense_shape.name], + [indices_np, values_np, dense_shape_np, default_value_int], + [indices.name, values.name, dense_shape.name, default_value.name], [ "sparse_fill_empty_rows/SparseFillEmptyRows:0", "sparse_fill_empty_rows/SparseFillEmptyRows:1", From 326ba03b067cdd05e3882a85e5b6c4e8391cbf58 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 03:58:14 +0000 Subject: [PATCH 08/32] Add Frontend Code --- .../frontend/tensorflow/test_forward.py | 18 ++++++----- .../relay/dyn/test_dynamic_op_level3.py | 30 +++++++++---------- 2 files changed, 26 insertions(+), 22 deletions(-) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 0fee664b7dfd..62d3107c7e38 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1824,7 +1824,7 @@ def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_ dense_shape = tf.placeholder( shape=dense_shape_np.shape, dtype=dense_shape_np.dtype, name="dense_shape" ) - default_value = tf.placeholder(shape=(), dtype = values_np.dtype, name="default_value") + default_value = tf.placeholder(shape=(), dtype=values_np.dtype, name="default_value") sp_input = tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=dense_shape) _ = tf.sparse.fill_empty_rows(sp_input, default_value, name="sparse_fill_empty_rows") compare_tf_with_tvm( @@ -1855,22 +1855,26 @@ def test_forward_sparse_fill_empty_rows(): sparse_indices_np = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64) sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int64) dense_shape_np = np.array([5, 6], dtype=np.int64) - default_value = 10 - _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, dense_shape_np, default_value) + default_value_int = 10 + _test_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int + ) sparse_indices_np = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64) sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int64) dense_shape_np = np.array([7, 7, 7], dtype=np.int64) - default_value_np = 5 + default_value_int = 5 _test_sparse_fill_empty_rows( - sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np + sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int ) sparse_indices_np = np.array([[1], [2]], dtype=np.int64) sparse_values_np = np.array([7, 8], dtype=np.int64) dense_shape_np = np.array([5], dtype=np.int64) - default_value_np = 4 - _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, dense_shape_np, default_value) + default_value_int = 4 + _test_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int + ) ####################################################################### diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index fa993e61b577..a82118df9094 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -278,28 +278,28 @@ def verify_sparse_fill_empty_rows( func, [sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np], ref_res ) - sparse_indices = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int32) - sparse_values = np.array([1, 2, 3, 4], dtype=np.int32) - dense_shape = np.array([7, 6], dtype=np.int32) + sparse_indices = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64) + sparse_values = np.array([1, 2, 3, 4], dtype=np.int64) + dense_shape = np.array([7, 6], dtype=np.int64) default_value = np.array([5], dtype=sparse_values.dtype) verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - sparse_indices = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int32) - sparse_values = np.array([1, 2, 3, 4], dtype=np.int32) - dense_shape = np.array([7, 7, 7], dtype=np.int32) - default_value = np.array([10], dtype=np.int32) + sparse_indices = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64) + sparse_values = np.array([1, 2, 3, 4], dtype=np.int64) + dense_shape = np.array([7, 7, 7], dtype=np.int64) + default_value = np.array([10], dtype=np.int64) verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - sparse_indices = np.array([[1], [2]], dtype=np.int32) - sparse_values = np.array([7, 8], dtype=np.int32) - dense_shape = np.array([5], dtype=np.int32) - default_value = np.array([4], dtype=np.int32) + sparse_indices = np.array([[1], [2]], dtype=np.int64) + sparse_values = np.array([7, 8], dtype=np.int64) + dense_shape = np.array([5], dtype=np.int64) + default_value = np.array([4], dtype=np.int64) verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - # sparse_indices = np.ones((0, 1), dtype=np.int32) - # sparse_values = np.array([], dtype=np.int32) - # dense_shape = np.array([5], dtype=np.int32) - # default_value = np.array([4], dtype=np.int32) + # sparse_indices = np.array([[]], dtype=np.int64) + # sparse_values = np.array([], dtype=np.int64) + # dense_shape = np.array([5], dtype=np.int64) + # default_value = np.array([4], dtype=np.int64) # verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) From 54d64cb7083c159fbe4349469a59fed362e566ca Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 10:57:52 +0000 Subject: [PATCH 09/32] New test Cases and new code to handle them --- python/tvm/relay/op/_transform.py | 42 +++++---- python/tvm/topi/sparse_fill_empty_rows.py | 85 +++++++++++-------- .../frontend/tensorflow/test_forward.py | 63 +++++++++----- .../relay/dyn/test_dynamic_op_level3.py | 16 ++-- 4 files changed, 125 insertions(+), 81 deletions(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 760b88891d50..caf93c7c4e1b 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -465,25 +465,35 @@ def argwhere_shape_func(attrs, inputs, out_ndims): @script def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): + new_sparse_indices_shape = output_tensor((2,), "int64") new_sparse_values_shape = output_tensor((1,), "int64") empty_row_indicator_shape = output_tensor((1,), "int64") - num_rows = int64(dense_shape[0]) - count = int64(sparse_indices.shape[0]) - for i in const_range(1, sparse_indices.shape[0]): - index = int64(sparse_indices[i, 0]) - prev_index = int64(sparse_indices[i - 1, 0] + 1) - - if index > prev_index: - count += index - prev_index - - count += int64(sparse_indices[0, 0]) - count += int64(num_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]) - new_sparse_indices_shape[0] = int64(count) - new_sparse_indices_shape[1] = int64(sparse_indices.shape[1]) - new_sparse_values_shape[0] = int64(count) - empty_row_indicator_shape[0] = int64(dense_shape[0]) - return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape) + num_dense_rows = int64(dense_shape[0]) + + if int64(sparse_indices.shape[0]) == int64(0): + new_sparse_indices_shape[0] = num_dense_rows + new_sparse_indices_shape[1] = int64(sparse_indices.shape[1]) + new_sparse_values_shape[0] = num_dense_rows + empty_row_indicator_shape[0] = num_dense_rows + return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape) + + else: + count = int64(sparse_indices.shape[0]) + for i in const_range(1, sparse_indices.shape[0]): + index = int64(sparse_indices[i, 0]) + prev_index = int64(sparse_indices[i - 1, 0] + 1) + + if index > prev_index: + count += index - prev_index + + count += int64(sparse_indices[0, 0]) + count += int64(num_dense_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]) + new_sparse_indices_shape[0] = int64(count) + new_sparse_indices_shape[1] = int64(sparse_indices.shape[1]) + new_sparse_values_shape[0] = int64(count) + empty_row_indicator_shape[0] = num_dense_rows + return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape) @_reg.register_shape_func("sparse_fill_empty_rows", True) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index 15955b1244f9..5f196df5efea 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -36,53 +36,64 @@ def _sparse_fill_empty_rows( empty_row_indicator = output_tensor(empty_row_indicator_shape, "int64") idx = 0 - for i in range(0, int64(sparse_indices[0, 0])): - new_sparse_indices[idx, 0] = int64(i) - for k in const_range(1, sparse_indices.shape[1]): - new_sparse_indices[idx, k] = int64(0) + if int64(sparse_indices.shape[0]) == int64(0): + for i in range(0, int64(new_sparse_indices_shape[0])): + new_sparse_indices[i, 0] = int64(i) + new_sparse_values[i] = default_value_ + empty_row_indicator[i] = int64(1) + for k in range(1, int64(new_sparse_indices_shape[1])): + new_sparse_indices[i, k] = int64(0) - new_sparse_values[idx] = default_value_ - empty_row_indicator[i] = int64(1) - idx += 1 + return (new_sparse_indices, new_sparse_values, empty_row_indicator) - for i in const_range(0, sparse_indices.shape[0]): - index = int64(sparse_indices[i, 0]) - if i == 0: - new_sparse_indices[idx, 0] = index + else: + for i in range(0, int64(sparse_indices[0, 0])): + new_sparse_indices[idx, 0] = int64(i) for k in const_range(1, sparse_indices.shape[1]): - new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) - new_sparse_values[idx] = int64(sparse_values[i]) - empty_row_indicator[index] = int64(0) + new_sparse_indices[idx, k] = int64(0) + + new_sparse_values[idx] = default_value_ + empty_row_indicator[i] = int64(1) idx += 1 - else: - prev_index = int64(sparse_indices[i - 1, 0] + 1) - for j in range(prev_index, index): - new_sparse_indices[idx, 0] = int64(j) + + for i in const_range(0, sparse_indices.shape[0]): + index = int64(sparse_indices[i, 0]) + if i == 0: + new_sparse_indices[idx, 0] = index for k in const_range(1, sparse_indices.shape[1]): - new_sparse_indices[idx, k] = int64(0) - empty_row_indicator[prev_index] = int64(1) - new_sparse_values[idx] = default_value_ + new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) + new_sparse_values[idx] = int64(sparse_values[i]) + empty_row_indicator[index] = int64(0) idx += 1 + else: + prev_index = int64(sparse_indices[i - 1, 0] + 1) + for j in range(prev_index, index): + new_sparse_indices[idx, 0] = int64(j) + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(0) + empty_row_indicator[prev_index] = int64(1) + new_sparse_values[idx] = default_value_ + idx += 1 - new_sparse_indices[idx, 0] = index - for k in const_range(1, sparse_indices.shape[1]): - new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) - new_sparse_values[idx] = int64(sparse_values[i]) - empty_row_indicator[index] = int64(0) - idx += 1 + new_sparse_indices[idx, 0] = index + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) + new_sparse_values[idx] = int64(sparse_values[i]) + empty_row_indicator[index] = int64(0) + idx += 1 - for i in range( - int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1), int64(dense_shape[0]) - ): + for i in range( + int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1), int64(dense_shape[0]) + ): - new_sparse_indices[idx, 0] = int64(i) - for k in const_range(1, sparse_indices.shape[1]): - new_sparse_indices[idx, k] = int64(0) - empty_row_indicator[i] = int64(1) - new_sparse_values[idx] = default_value_ - idx += 1 + new_sparse_indices[idx, 0] = int64(i) + for k in const_range(1, sparse_indices.shape[1]): + new_sparse_indices[idx, k] = int64(0) + empty_row_indicator[i] = int64(1) + new_sparse_values[idx] = default_value_ + idx += 1 - return (new_sparse_indices, new_sparse_values, empty_row_indicator) + return (new_sparse_indices, new_sparse_values, empty_row_indicator) def sparse_fill_empty_rows( diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 62d3107c7e38..94d6a9cdcdee 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -254,6 +254,7 @@ def name_without_num(name): ) # since the names from tensorflow and relay runs are not exactly same, # first len(tf_output) will be compared + print(tvm_output) for i in range(len(tf_output)): if not isinstance(tf_output[i], np.ndarray): assert len(tvm_output[i].shape) == 0 @@ -1839,7 +1840,44 @@ def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_ ) -def test_forward_sparse_fill_empty_rows(): +@pytest.mark.parametrize( + "sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int", + [ + ( + np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64), + np.array([1, 2, 3, 4], dtype=np.int64), + np.array([5, 6], dtype=np.int64), + 10, + ), + ( + np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64), + np.array([1, 2, 3, 4], dtype=np.int64), + np.array([7, 7, 7], dtype=np.int64), + 5, + ), + ( + np.array([[1], [2]], dtype=np.int64), + np.array([7, 8], dtype=np.int64), + np.array([5], dtype=np.int64), + 4, + ), + ( + np.ones((0, 1), dtype=np.int64), + np.array([], dtype=np.int64), + np.array([5], dtype=np.int64), + 4, + ), + ( + np.ones((0, 3), dtype=np.int64), + np.array([], dtype=np.int64), + np.array([9, 3, 7], dtype=np.int64), + 100, + ), + ], +) +def test_forward_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int +): """ sparse_fill_empty_rows op test""" ################################################################### # @@ -1852,26 +1890,6 @@ def test_forward_sparse_fill_empty_rows(): # [0, 0, 0, 0]] # # ------------------------------------------------------------------ - sparse_indices_np = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64) - sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int64) - dense_shape_np = np.array([5, 6], dtype=np.int64) - default_value_int = 10 - _test_sparse_fill_empty_rows( - sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int - ) - - sparse_indices_np = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64) - sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int64) - dense_shape_np = np.array([7, 7, 7], dtype=np.int64) - default_value_int = 5 - _test_sparse_fill_empty_rows( - sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int - ) - - sparse_indices_np = np.array([[1], [2]], dtype=np.int64) - sparse_values_np = np.array([7, 8], dtype=np.int64) - dense_shape_np = np.array([5], dtype=np.int64) - default_value_int = 4 _test_sparse_fill_empty_rows( sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int ) @@ -4752,5 +4770,4 @@ def lstm_cell(): if __name__ == "__main__": - test_forward_sparse_fill_empty_rows() - # pytest.main([__file__]) + pytest.main([__file__]) diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index a82118df9094..6d0a84da4bbe 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -296,11 +296,17 @@ def verify_sparse_fill_empty_rows( default_value = np.array([4], dtype=np.int64) verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - # sparse_indices = np.array([[]], dtype=np.int64) - # sparse_values = np.array([], dtype=np.int64) - # dense_shape = np.array([5], dtype=np.int64) - # default_value = np.array([4], dtype=np.int64) - # verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) + sparse_indices = np.ones((0, 1), dtype=np.int64) + sparse_values = np.array([], dtype=np.int64) + dense_shape = np.array([5], dtype=np.int64) + default_value = np.array([4], dtype=np.int64) + verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) + + sparse_indices = np.ones((0, 3), dtype=np.int64) + sparse_values = np.array([], dtype=np.int64) + dense_shape = np.array([9, 3, 7], dtype=np.int64) + default_value = np.array([100], dtype=np.int64) + verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) if __name__ == "__main__": From af86e6f3e9a41ea636355f9bf682317dd7dc784b Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 12:06:51 +0000 Subject: [PATCH 10/32] Add Python Implementation' ' --- .../relay/dyn/test_dynamic_op_level3.py | 124 ++++++++++-------- 1 file changed, 68 insertions(+), 56 deletions(-) diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 6d0a84da4bbe..1acb62412fb2 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -29,14 +29,18 @@ def verify_func(func, data, ref_res): assert isinstance(data, list) for target, ctx in tvm.testing.enabled_targets(): - for kind in ["vm"]: + for kind in ["vm", "debug"]: mod = tvm.ir.IRModule.from_expr(func) intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target) op_res = intrp.evaluate()(*data) - for op_result, ref_result in zip(op_res, ref_res): - # print(op_result.asnumpy(), ref_result) - print(op_result.asnumpy()) - # tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5) + if isinstance(op_res, tvm.runtime.container.ADT): + assert len(op_res) == len( + ref_res + ), "Outputs from TVM and Python implementation must be equal " + for op_result, ref_result in zip(op_res, ref_res): + tvm.testing.assert_allclose(op_result.asnumpy(), ref_result, rtol=1e-5) + else: + tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5) relay.backend.compile_engine.get().clear() @@ -205,7 +209,42 @@ def verify_sparse_to_dense(sparse_indices, sparse_values, default_value, output_ verify_sparse_to_dense(1, 3, None, [5], [0, 3, 0, 0, 0]) # default value not specified -def test_sparse_fill_empty_rows(): +@pytest.mark.parametrize( + "sparse_indices, sparse_values, dense_shape, default_value", + [ + ( + np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64), + np.array([1, 2, 3, 4], dtype=np.int64), + np.array([5, 6], dtype=np.int64), + np.array([10], dtype=np.int64), + ), + ( + np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64), + np.array([1, 2, 3, 4], dtype=np.int64), + np.array([7, 7, 7], dtype=np.int64), + np.array([5], dtype=np.int64), + ), + ( + np.array([[1], [2]], dtype=np.int64), + np.array([7, 8], dtype=np.int64), + np.array([5], dtype=np.int64), + np.array([4], dtype=np.int64), + ), + ( + np.ones((0, 1), dtype=np.int64), + np.array([], dtype=np.int64), + np.array([5], dtype=np.int64), + np.array([4], dtype=np.int64), + ), + ( + np.ones((0, 3), dtype=np.int64), + np.array([], dtype=np.int64), + np.array([9, 3, 7], dtype=np.int64), + np.array([100], dtype=np.int64), + ), + ], +) +def test_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): def ref_sparse_fill_empty_rows( sparse_indices: np.ndarray, sparse_values: np.ndarray, @@ -216,26 +255,29 @@ def ref_sparse_fill_empty_rows( This function calculates the expected output of sparse_fill_empty_rows operator given the inputs. """ - new_sparse_indices = -1 * np.ones( - (sparse_indices.shape[0] + dense_shape[0], sparse_indices.shape[1]) - ) - empty_row_indicator = np.ones(dense_shape[0], dtype=bool) - new_sparse_values = -1 * np.ones(sparse_indices.shape[0] + dense_shape[0]) - - for i in range(sparse_indices.shape[0]): - empty_row_indicator[sparse_indices[i][0]] = False - - new_sparse_indices[: sparse_indices.shape[0]][:] = sparse_indices[:] - new_sparse_values[: sparse_values.shape[0]] = sparse_values[:] - new_sparse_indices_index = sparse_indices.shape[0] - - for empty_row_index, element in enumerate(empty_row_indicator): - if element: - new_sparse_indices[new_sparse_indices_index, 0] = empty_row_index - new_sparse_indices[new_sparse_indices_index, 1:] = 0 - new_sparse_values[new_sparse_indices_index] = default_value[0] - new_sparse_indices_index += 1 + def check_add_rows(current_idx, limit_idx): + while current_idx < limit_idx: + new_sparse_indices.append([current_idx] + [0] * (num_cols - 1)) + new_sparse_values.append(default_value[0]) + empty_row_indicator[current_idx] = 1 + current_idx += 1 + + return current_idx + + current_idx = 0 + new_sparse_indices = [] + new_sparse_values = [] + empty_row_indicator = [0 for _ in range(dense_shape[0])] + num_cols = sparse_indices.shape[1] + for sparse_row, sparse_value in zip(sparse_indices, sparse_values): + limit_idx = sparse_row[0] + current_idx = check_add_rows(current_idx, limit_idx) + new_sparse_indices.append(list(sparse_row)) + new_sparse_values.append(sparse_value) + current_idx = limit_idx + 1 + + check_add_rows(current_idx, dense_shape[0]) return new_sparse_indices, new_sparse_values, empty_row_indicator def verify_sparse_fill_empty_rows( @@ -273,42 +315,12 @@ def verify_sparse_fill_empty_rows( dense_shape_np, default_value_np, ) - # ref_res = [] verify_func( func, [sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np], ref_res ) - sparse_indices = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64) - sparse_values = np.array([1, 2, 3, 4], dtype=np.int64) - dense_shape = np.array([7, 6], dtype=np.int64) - default_value = np.array([5], dtype=sparse_values.dtype) - verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - - sparse_indices = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int64) - sparse_values = np.array([1, 2, 3, 4], dtype=np.int64) - dense_shape = np.array([7, 7, 7], dtype=np.int64) - default_value = np.array([10], dtype=np.int64) - verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - - sparse_indices = np.array([[1], [2]], dtype=np.int64) - sparse_values = np.array([7, 8], dtype=np.int64) - dense_shape = np.array([5], dtype=np.int64) - default_value = np.array([4], dtype=np.int64) - verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - - sparse_indices = np.ones((0, 1), dtype=np.int64) - sparse_values = np.array([], dtype=np.int64) - dense_shape = np.array([5], dtype=np.int64) - default_value = np.array([4], dtype=np.int64) - verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) - - sparse_indices = np.ones((0, 3), dtype=np.int64) - sparse_values = np.array([], dtype=np.int64) - dense_shape = np.array([9, 3, 7], dtype=np.int64) - default_value = np.array([100], dtype=np.int64) verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) if __name__ == "__main__": - test_sparse_fill_empty_rows() - # pytest.main([__file__]) + pytest.main([__file__]) From 7281a759c571a8043d09c76dcb57b2b43531d61b Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 12:08:14 +0000 Subject: [PATCH 11/32] Remove stuff --- tests/python/frontend/tensorflow/test_forward.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 94d6a9cdcdee..ef5fdeeca1b8 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -254,7 +254,6 @@ def name_without_num(name): ) # since the names from tensorflow and relay runs are not exactly same, # first len(tf_output) will be compared - print(tvm_output) for i in range(len(tf_output)): if not isinstance(tf_output[i], np.ndarray): assert len(tvm_output[i].shape) == 0 From d432a192b94f5d04af632099e59e3b81c3f6f84c Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 12:13:58 +0000 Subject: [PATCH 12/32] Remove unused imports --- python/tvm/topi/sparse_fill_empty_rows.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index 5f196df5efea..c3ab245d6484 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -16,8 +16,7 @@ # under the License. # pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks """SparseFillEmptyRows operator""" -from ..tir import decl_buffer, ir_builder, Cast, AssertStmt, StringImm, Evaluate -from ..te import extern, hybrid +from ..te import hybrid @hybrid.script From 7dff190cfdfc3660aa8d2b87879f6fdfe6d1fef2 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 12:17:48 +0000 Subject: [PATCH 13/32] Pylint --- python/tvm/topi/sparse_fill_empty_rows.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index c3ab245d6484..aa3d4575549a 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -14,7 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks + +# pylint: no-else-return """SparseFillEmptyRows operator""" from ..te import hybrid From 58df4a8f5302a860cfbda9f82a5640834786fe2a Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 12:28:04 +0000 Subject: [PATCH 14/32] Pylint --- python/tvm/topi/sparse_fill_empty_rows.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index aa3d4575549a..8134be92e36f 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. -# pylint: no-else-return +# pylint: disable=no-else-return, too-many-locals, too-many-arguments, too-many-branches +# pylint: disable=undefined-variable, invalid-name """SparseFillEmptyRows operator""" from ..te import hybrid From 701328935e5f04b9e59aa595eb0c7135e5b3adf2 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 12:33:48 +0000 Subject: [PATCH 15/32] PyLint Shape Func --- python/tvm/relay/op/_transform.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index caf93c7c4e1b..1ed5d595e01a 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -15,7 +15,9 @@ # specific language governing permissions and limitations # under the License. """Backend compiler related feature registration""" -# pylint: disable=invalid-name,unused-argument, len-as-condition, too-many-nested-blocks, too-many-local-variables, too-many-arguments +# pylint: disable=invalid-name,unused-argument, len-as-condition, too-many-nested-blocks, +# pylint: disable=too-many-local-variables, too-many-arguments, no-else-return + from __future__ import absolute_import import tvm from tvm import te From 1c83857fc5c32c38020ad8347bde5f68bf8bb3c7 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 12:49:49 +0000 Subject: [PATCH 16/32] Make tests cpu only --- tests/python/relay/dyn/test_dynamic_op_level3.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 1acb62412fb2..95b607652e8a 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -26,9 +26,9 @@ import tvm.testing -def verify_func(func, data, ref_res): +def verify_func(func, data, ref_res, target_ctx=tvm.testing.enabled_targets()): assert isinstance(data, list) - for target, ctx in tvm.testing.enabled_targets(): + for target, ctx in target_ctx: for kind in ["vm", "debug"]: mod = tvm.ir.IRModule.from_expr(func) intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target) @@ -316,7 +316,10 @@ def verify_sparse_fill_empty_rows( default_value_np, ) verify_func( - func, [sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np], ref_res + func, + [sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np], + ref_res, + [("llvm", tvm.cpu())], ) verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) From b6d3603272a9cf0edd951d30f932be06e263f538 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 22:43:28 +0000 Subject: [PATCH 17/32] Add unsorted tests --- tests/python/frontend/tensorflow/test_forward.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index ef5fdeeca1b8..0582c054f2f4 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1842,6 +1842,12 @@ def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_ @pytest.mark.parametrize( "sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int", [ + ( + np.array([[1, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64), + np.array([1, 2, 3, 4], dtype=np.int64), + np.array([5, 6], dtype=np.int64), + 10, + ), ( np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64), np.array([1, 2, 3, 4], dtype=np.int64), From f3fdc0a5623fc6b529ac0b0f2f5ee8b56f1fd880 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 22:43:49 +0000 Subject: [PATCH 18/32] Add frontend code --- python/tvm/relay/frontend/tensorflow.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index fd6505d3abdf..6d6d452d662c 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -1001,8 +1001,15 @@ def _impl(inputs, attr, params, mod): def _sparse_fill_empty_rows(): def _impl(inputs, attr, params, mod): assert len(inputs) == 4, "There should be 4 input tensors" + sparse_indices = inputs[0] + sparse_values = inputs[1] + sparse_indices_num_cols = _infer_shape(sparse_indices, mod)[1] + first_column = _op.split(sparse_indices, sparse_indices_num_cols, axis=1)[0] + sorted_indices = _op.argsort(_op.squeeze(first_column)) + sorted_sparse_indices = _op.take(sparse_indices, sorted_indices, axis=0) + sorted_sparse_values = _op.take(sparse_values, sorted_indices, axis=0) new_sparse_indices, new_sparse_values, empty_row_indicator = _op.sparse_fill_empty_rows( - inputs[0], inputs[1], inputs[2], inputs[3] + sorted_sparse_indices, sorted_sparse_values, inputs[2], inputs[3] ) return _expr.TupleWrapper( From 07062f6c4f4bf61e7730d69880e5bd6db5d28056 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 11 Feb 2021 22:47:35 +0000 Subject: [PATCH 19/32] Row Major Sorting Only Test --- tests/python/frontend/tensorflow/test_forward.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 0582c054f2f4..2538f932cbda 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1842,6 +1842,12 @@ def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_ @pytest.mark.parametrize( "sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int", [ + ( + np.array([[1, 1], [0, 3], [0, 1], [2, 0], [3, 1]], dtype=np.int64), + np.array([1, 2, 3, 4, 5], dtype=np.int64), + np.array([5, 6], dtype=np.int64), + 10, + ), ( np.array([[1, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int64), np.array([1, 2, 3, 4], dtype=np.int64), From 83164946e0df3113202961d8a9d02ac1bff53739 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Fri, 12 Feb 2021 21:25:57 +0000 Subject: [PATCH 20/32] Handle Dynamic Shapes --- python/tvm/relay/op/_transform.py | 2 +- python/tvm/topi/sparse_fill_empty_rows.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 1ed5d595e01a..c43718823c7d 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -482,7 +482,7 @@ def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): else: count = int64(sparse_indices.shape[0]) - for i in const_range(1, sparse_indices.shape[0]): + for i in range(1, int64(sparse_indices.shape[0])): index = int64(sparse_indices[i, 0]) prev_index = int64(sparse_indices[i - 1, 0] + 1) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index 8134be92e36f..746fb661575b 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -50,18 +50,18 @@ def _sparse_fill_empty_rows( else: for i in range(0, int64(sparse_indices[0, 0])): new_sparse_indices[idx, 0] = int64(i) - for k in const_range(1, sparse_indices.shape[1]): + for k in range(1, int64(new_sparse_indices_shape[1])): new_sparse_indices[idx, k] = int64(0) new_sparse_values[idx] = default_value_ empty_row_indicator[i] = int64(1) idx += 1 - for i in const_range(0, sparse_indices.shape[0]): + for i in range(0, int64(sparse_indices.shape[0])): index = int64(sparse_indices[i, 0]) if i == 0: new_sparse_indices[idx, 0] = index - for k in const_range(1, sparse_indices.shape[1]): + for k in range(1, int64(new_sparse_indices_shape[1])): new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) new_sparse_values[idx] = int64(sparse_values[i]) empty_row_indicator[index] = int64(0) @@ -70,14 +70,14 @@ def _sparse_fill_empty_rows( prev_index = int64(sparse_indices[i - 1, 0] + 1) for j in range(prev_index, index): new_sparse_indices[idx, 0] = int64(j) - for k in const_range(1, sparse_indices.shape[1]): + for k in range(1, int64(new_sparse_indices_shape[1])): new_sparse_indices[idx, k] = int64(0) empty_row_indicator[prev_index] = int64(1) new_sparse_values[idx] = default_value_ idx += 1 new_sparse_indices[idx, 0] = index - for k in const_range(1, sparse_indices.shape[1]): + for k in range(1, int64(new_sparse_indices_shape[1])): new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) new_sparse_values[idx] = int64(sparse_values[i]) empty_row_indicator[index] = int64(0) @@ -88,7 +88,7 @@ def _sparse_fill_empty_rows( ): new_sparse_indices[idx, 0] = int64(i) - for k in const_range(1, sparse_indices.shape[1]): + for k in range(1, int64(new_sparse_indices_shape[1])): new_sparse_indices[idx, k] = int64(0) empty_row_indicator[i] = int64(1) new_sparse_values[idx] = default_value_ From e2781c14281a15b9b494f4748d7523bfb66dd0f4 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Sat, 13 Feb 2021 02:01:28 +0000 Subject: [PATCH 21/32] Add dynamic input shapes --- python/tvm/relay/op/transform.py | 8 ++++++ python/tvm/topi/transform.py | 8 ++++++ .../frontend/tensorflow/test_forward.py | 25 +++++++++++++------ 3 files changed, 33 insertions(+), 8 deletions(-) diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 20e1d04f99d7..684891af6244 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1323,6 +1323,14 @@ def adv_index(inputs): def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): + """ + This op exactly follows the documentation here: + https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows + There are two exceptions: + 1. Input Sparse Indices are expected to be in row-major order. + 2. Empty Row Indicator has int64 output type with 1(for True) and 0(for False). + + """ return TupleWrapper( _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 3 ) diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index 0ac9bf9f9715..eb026aafffb0 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -934,4 +934,12 @@ def adv_index(data, indices): def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): + """ + This op exactly follows the documentation here: + https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows + There are two exceptions: + 1. Input Sparse Indices are expected to be in row-major order. + 2. Empty Row Indicator has int64 output type with 1(for True) and 0(for False). + + """ return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 2538f932cbda..7ecdb7882502 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1817,13 +1817,21 @@ def test_forward_sparse_dense_matmul(): # ------------ -def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_value_int): +def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_value_int, use_dyn): with tf.Graph().as_default(): - indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices") - values = tf.placeholder(shape=values_np.shape, dtype=values_np.dtype, name="values") - dense_shape = tf.placeholder( - shape=dense_shape_np.shape, dtype=dense_shape_np.dtype, name="dense_shape" - ) + if use_dyn: + indices = tf.placeholder(shape=(None, None), dtype=indices_np.dtype, name="indices") + values = tf.placeholder(shape=(None), dtype=values_np.dtype, name="values") + dense_shape = tf.placeholder( + shape=(None), dtype=dense_shape_np.dtype, name="dense_shape" + ) + else: + indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices") + values = tf.placeholder(shape=values_np.shape, dtype=values_np.dtype, name="values") + dense_shape = tf.placeholder( + shape=dense_shape_np.shape, dtype=dense_shape_np.dtype, name="dense_shape" + ) + default_value = tf.placeholder(shape=(), dtype=values_np.dtype, name="default_value") sp_input = tf.sparse.SparseTensor(indices=indices, values=values, dense_shape=dense_shape) _ = tf.sparse.fill_empty_rows(sp_input, default_value, name="sparse_fill_empty_rows") @@ -1886,8 +1894,9 @@ def _test_sparse_fill_empty_rows(indices_np, values_np, dense_shape_np, default_ ), ], ) +@pytest.mark.parametrize("use_dyn", [True, False]) def test_forward_sparse_fill_empty_rows( - sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int + sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int, use_dyn ): """ sparse_fill_empty_rows op test""" ################################################################### @@ -1902,7 +1911,7 @@ def test_forward_sparse_fill_empty_rows( # # ------------------------------------------------------------------ _test_sparse_fill_empty_rows( - sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int + sparse_indices_np, sparse_values_np, dense_shape_np, default_value_int, use_dyn ) From 856bd6c0bb1fa1da8479b675a94b6fac1094db46 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Sat, 13 Feb 2021 02:02:40 +0000 Subject: [PATCH 22/32] Dynamic Shape Tests --- .../relay/dyn/test_dynamic_op_level3.py | 57 +++++++++++++------ 1 file changed, 40 insertions(+), 17 deletions(-) diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 95b607652e8a..12b15f06e976 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -244,7 +244,8 @@ def verify_sparse_to_dense(sparse_indices, sparse_values, default_value, output_ ), ], ) -def test_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): +@pytest.mark.parametrize("use_dyn", [True, False]) +def test_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value, use_dyn): def ref_sparse_fill_empty_rows( sparse_indices: np.ndarray, sparse_values: np.ndarray, @@ -289,22 +290,44 @@ def verify_sparse_fill_empty_rows( """ This function verifies the relay output of sparse_fill_empty_rows with its expected output. """ - sparse_indices = relay.var( - "sparse_indices", - relay.TensorType(sparse_indices_np.shape, str(sparse_indices_np.dtype)), - ) - sparse_values = relay.var( - "sparse_values", - relay.TensorType(sparse_values_np.shape, str(sparse_values_np.dtype)), - ) - dense_shape = relay.var( - "dense_shape", - relay.TensorType(dense_shape_np.shape, str(dense_shape_np.dtype)), - ) - default_value = relay.var( - "default_value", - relay.TensorType(default_value_np.shape, str(default_value_np.dtype)), - ) + if use_dyn: + sparse_indices = relay.var( + "sparse_indices", + shape=[relay.Any(), relay.Any()], + dtype=str(sparse_indices_np.dtype), + ) + sparse_values = relay.var( + "sparse_values", + shape=[relay.Any()], + dtype=str(sparse_values_np.dtype), + ) + dense_shape = relay.var( + "dense_shape", + shape=[relay.Any()], + dtype=str(dense_shape_np.dtype), + ) + default_value = relay.var( + "default_value", + shape=[relay.Any()], + dtype=str(default_value_np.dtype), + ) + else: + sparse_indices = relay.var( + "sparse_indices", + relay.TensorType(sparse_indices_np.shape, str(sparse_indices_np.dtype)), + ) + sparse_values = relay.var( + "sparse_values", + relay.TensorType(sparse_values_np.shape, str(sparse_values_np.dtype)), + ) + dense_shape = relay.var( + "dense_shape", + relay.TensorType(dense_shape_np.shape, str(dense_shape_np.dtype)), + ) + default_value = relay.var( + "default_value", + relay.TensorType(default_value_np.shape, str(default_value_np.dtype)), + ) z = relay.sparse_fill_empty_rows( sparse_indices, sparse_values, dense_shape, default_value ).astuple() From 008d8d9c0ade3cb9993b4b361ae7b80d0c9b6617 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 15 Feb 2021 20:37:25 +0000 Subject: [PATCH 23/32] Add documentation --- python/tvm/relay/op/transform.py | 51 ++++++++++++++++++++++++++++++++ python/tvm/topi/transform.py | 51 ++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 684891af6244..245b4ec4e77d 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1324,12 +1324,63 @@ def adv_index(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 3 outputs + Parameters + ---------- + sparse_indices : relay.Expr + A 2-D int64 tensor[N, ndims] 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. + The inputs for this relay parameter must be in row-major order. + sparse_values : relay.Expr + A 1-D int64 tensor[N] containing the sparse values for the sparse indices. + dense_shape : relay.Expr + A 1-D int64 tensor[ndims] which contains shape of the dense output tensor. + default_value : relay.Expr + A 1-D tensor[1] containing the default value for the remaining locations. + Returns + ------- + new_sparse_indices : relay.Expr + A 2-D int64 tensor[?, ndims] of integers containing location of new sparse + indices + new_sparse_values : relay.Expr + A 1-D int64 tensor[?] containing the sparse values for the sparse indices. + empty_row_indicator : relay.Expr + A 1-D int64 tensor[dense_shape[0]] filled with zeros and ones + indicating whether the particular row is empty or full respectively + + Note: This op exactly follows the documentation here: https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows There are two exceptions: 1. Input Sparse Indices are expected to be in row-major order. 2. Empty Row Indicator has int64 output type with 1(for True) and 0(for False). + Examples + ------- + .. code-block:: python + sparse_indices = [[0, 1], + [0, 3], + [2, 0], + [3, 1]] + sparse_values = [1, 2, 3, 4] + default_value = [10] + dense_shape = [5, 6] + new_sparse_indices, empty_row_indicator, new_sparse_values, slice_element_index = + relay.sparse_fill_empty_rows( + sparse_indices, + sparse_values, + default_value, + dense_shape) + new_sparse_indices = [[0, 1], + [0, 3], + [1, 0], + [2, 0], + [3, 1], + [4, 0]] + empty_row_indicator = [0, 1, 0, 0, 1] + new_sparse_values = [1, 2, 10, 3, 4, 10] + """ return TupleWrapper( _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 3 diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index eb026aafffb0..794fa53c786f 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -935,11 +935,62 @@ def adv_index(data, indices): 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 3 outputs + Parameters + ---------- + sparse_indices : relay.Expr + A 2-D int64 tensor[N, ndims] 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. + The inputs for this relay parameter must be in row-major order. + sparse_values : relay.Expr + A 1-D int64 tensor[N] containing the sparse values for the sparse indices. + dense_shape : relay.Expr + A 1-D int64 tensor[ndims] which contains shape of the dense output tensor. + default_value : relay.Expr + A 1-D tensor[1] containing the default value for the remaining locations. + Returns + ------- + new_sparse_indices : relay.Expr + A 2-D int64 tensor[?, ndims] of integers containing location of new sparse + indices + new_sparse_values : relay.Expr + A 1-D int64 tensor[?] containing the sparse values for the sparse indices. + empty_row_indicator : relay.Expr + A 1-D int64 tensor[dense_shape[0]] filled with zeros and ones + indicating whether the particular row is empty or full respectively + + Note: This op exactly follows the documentation here: https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows There are two exceptions: 1. Input Sparse Indices are expected to be in row-major order. 2. Empty Row Indicator has int64 output type with 1(for True) and 0(for False). + Examples + ------- + .. code-block:: python + sparse_indices = [[0, 1], + [0, 3], + [2, 0], + [3, 1]] + sparse_values = [1, 2, 3, 4] + default_value = [10] + dense_shape = [5, 6] + new_sparse_indices, empty_row_indicator, new_sparse_values, slice_element_index = + relay.sparse_fill_empty_rows( + sparse_indices, + sparse_values, + default_value, + dense_shape) + new_sparse_indices = [[0, 1], + [0, 3], + [1, 0], + [2, 0], + [3, 1], + [4, 0]] + empty_row_indicator = [0, 1, 0, 0, 1] + new_sparse_values = [1, 2, 10, 3, 4, 10] + """ return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) From bdf9ab2c2feb414a7fc26574005817ad15d1d7e9 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 15 Feb 2021 20:42:15 +0000 Subject: [PATCH 24/32] Dtypes --- python/tvm/topi/sparse_fill_empty_rows.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index 746fb661575b..b5f8a5876db8 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -32,8 +32,8 @@ def _sparse_fill_empty_rows( empty_row_indicator_shape, ): default_value_ = int64(default_value[0]) - new_sparse_indices = output_tensor(new_sparse_indices_shape, "int64") - new_sparse_values = output_tensor(new_sparse_values_shape, "int64") + new_sparse_indices = output_tensor(new_sparse_indices_shape, sparse_indices.dtype) + new_sparse_values = output_tensor(new_sparse_values_shape, sparse_values.dtype) empty_row_indicator = output_tensor(empty_row_indicator_shape, "int64") idx = 0 From 23d3f4492c77d3efbc75427f6afa690a26bb89a6 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 15 Feb 2021 20:48:55 +0000 Subject: [PATCH 25/32] PR Comments --- src/relay/op/tensor/transform.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 02e4a5e1e113..ae6cebd99878 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1608,8 +1608,8 @@ 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) + .describe( + R"code(Fill empty rows of a sparse tensor with a default value.)code" TVM_ADD_FILELINE) .set_num_inputs(4) .add_argument("sparse_indices", "Tensor", "A 2-D int64 tensor of shape [N, ndims], which specifies the indices of the" From 666eb714dbc302d1c6f1470640153271353d8a15 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 15 Feb 2021 21:05:36 +0000 Subject: [PATCH 26/32] Added comments and changed naming --- python/tvm/topi/sparse_fill_empty_rows.py | 62 +++++++++++++---------- 1 file changed, 34 insertions(+), 28 deletions(-) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index b5f8a5876db8..791959640819 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -35,9 +35,10 @@ def _sparse_fill_empty_rows( new_sparse_indices = output_tensor(new_sparse_indices_shape, sparse_indices.dtype) new_sparse_values = output_tensor(new_sparse_values_shape, sparse_values.dtype) empty_row_indicator = output_tensor(empty_row_indicator_shape, "int64") - idx = 0 + new_sparse_indices_row_id = 0 - if int64(sparse_indices.shape[0]) == int64(0): + if int64(sparse_indices.shape[0]) == int64(0): # Handle Empty Case + # Fill all rows with default values for i in range(0, int64(new_sparse_indices_shape[0])): new_sparse_indices[i, 0] = int64(i) new_sparse_values[i] = default_value_ @@ -48,51 +49,56 @@ def _sparse_fill_empty_rows( return (new_sparse_indices, new_sparse_values, empty_row_indicator) else: + # Add rows with default value if first row id of sparse_indices is not a zero. for i in range(0, int64(sparse_indices[0, 0])): - new_sparse_indices[idx, 0] = int64(i) + new_sparse_indices[new_sparse_indices_row_id, 0] = int64(i) for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[idx, k] = int64(0) + new_sparse_indices[new_sparse_indices_row_id, k] = int64(0) - new_sparse_values[idx] = default_value_ + new_sparse_values[new_sparse_indices_row_id] = default_value_ empty_row_indicator[i] = int64(1) - idx += 1 + new_sparse_indices_row_id += 1 for i in range(0, int64(sparse_indices.shape[0])): - index = int64(sparse_indices[i, 0]) + row_id = int64(sparse_indices[i, 0]) if i == 0: - new_sparse_indices[idx, 0] = index + # Add first row of input to output + new_sparse_indices[new_sparse_indices_row_id, 0] = row_id for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) - new_sparse_values[idx] = int64(sparse_values[i]) - empty_row_indicator[index] = int64(0) - idx += 1 + new_sparse_indices[new_sparse_indices_row_id, k] = int64(sparse_indices[i, k]) + new_sparse_values[new_sparse_indices_row_id] = int64(sparse_values[i]) + empty_row_indicator[row_id] = int64(0) + new_sparse_indices_row_id += 1 else: - prev_index = int64(sparse_indices[i - 1, 0] + 1) - for j in range(prev_index, index): - new_sparse_indices[idx, 0] = int64(j) + prev_row_id = int64(sparse_indices[i - 1, 0] + 1) + # Since input is in row-major order, add rows between prev_row_id and row_id + for j in range(prev_row_id, row_id): + new_sparse_indices[new_sparse_indices_row_id, 0] = int64(j) for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[idx, k] = int64(0) - empty_row_indicator[prev_index] = int64(1) - new_sparse_values[idx] = default_value_ - idx += 1 + new_sparse_indices[new_sparse_indices_row_id, k] = int64(0) + empty_row_indicator[prev_row_id] = int64(1) + new_sparse_values[new_sparse_indices_row_id] = default_value_ + new_sparse_indices_row_id += 1 - new_sparse_indices[idx, 0] = index + # Add current row to output + new_sparse_indices[new_sparse_indices_row_id, 0] = row_id for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[idx, k] = int64(sparse_indices[i, k]) - new_sparse_values[idx] = int64(sparse_values[i]) - empty_row_indicator[index] = int64(0) - idx += 1 + new_sparse_indices[new_sparse_indices_row_id, k] = int64(sparse_indices[i, k]) + new_sparse_values[new_sparse_indices_row_id] = int64(sparse_values[i]) + empty_row_indicator[row_id] = int64(0) + new_sparse_indices_row_id += 1 + # Add rows with default value if last row id of sparse_indices is not dense_shape[0] - 1 for i in range( int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1), int64(dense_shape[0]) ): - new_sparse_indices[idx, 0] = int64(i) + new_sparse_indices[new_sparse_indices_row_id, 0] = int64(i) for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[idx, k] = int64(0) + new_sparse_indices[new_sparse_indices_row_id, k] = int64(0) empty_row_indicator[i] = int64(1) - new_sparse_values[idx] = default_value_ - idx += 1 + new_sparse_values[new_sparse_indices_row_id] = default_value_ + new_sparse_indices_row_id += 1 return (new_sparse_indices, new_sparse_values, empty_row_indicator) From d6b896494fa65ab2b6b1db0bd718d4e9a422d1c8 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 15 Feb 2021 21:07:50 +0000 Subject: [PATCH 27/32] Add comments --- python/tvm/topi/sparse_fill_empty_rows.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index 791959640819..8aa5f9b42314 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -49,7 +49,7 @@ def _sparse_fill_empty_rows( return (new_sparse_indices, new_sparse_values, empty_row_indicator) else: - # Add rows with default value if first row id of sparse_indices is not a zero. + # Add rows with default value if first row id of sparse_indices is not a zero for i in range(0, int64(sparse_indices[0, 0])): new_sparse_indices[new_sparse_indices_row_id, 0] = int64(i) for k in range(1, int64(new_sparse_indices_shape[1])): @@ -59,6 +59,7 @@ def _sparse_fill_empty_rows( empty_row_indicator[i] = int64(1) new_sparse_indices_row_id += 1 + # Iterate through sparse_indices and add rows if/when required for i in range(0, int64(sparse_indices.shape[0])): row_id = int64(sparse_indices[i, 0]) if i == 0: From f1db781b0b6558e609136097cf9338ca1d8b616e Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 15 Feb 2021 21:11:35 +0000 Subject: [PATCH 28/32] Comments to Shape Func --- python/tvm/relay/op/_transform.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index c43718823c7d..01bcf4a6cf60 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -473,7 +473,8 @@ def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): empty_row_indicator_shape = output_tensor((1,), "int64") num_dense_rows = int64(dense_shape[0]) - if int64(sparse_indices.shape[0]) == int64(0): + if int64(sparse_indices.shape[0]) == int64(0): # Handle Empty Case + # Total rows will equal dense_shape[0] new_sparse_indices_shape[0] = num_dense_rows new_sparse_indices_shape[1] = int64(sparse_indices.shape[1]) new_sparse_values_shape[0] = num_dense_rows @@ -481,16 +482,18 @@ def _sparse_fill_empty_rows_shape_func(sparse_indices, dense_shape): return (new_sparse_indices_shape, new_sparse_values_shape, empty_row_indicator_shape) else: - count = int64(sparse_indices.shape[0]) + count = int64(sparse_indices.shape[0]) # Add count of all rows already in sparse_indices for i in range(1, int64(sparse_indices.shape[0])): index = int64(sparse_indices[i, 0]) prev_index = int64(sparse_indices[i - 1, 0] + 1) if index > prev_index: - count += index - prev_index + count += index - prev_index # Add count of all rows between two consecutive indices - count += int64(sparse_indices[0, 0]) - count += int64(num_dense_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0]) + count += int64(sparse_indices[0, 0]) # Add count from 0 to first row id in sparse_indices + count += int64( + num_dense_rows - 1 - sparse_indices[sparse_indices.shape[0] - 1, 0] + ) # Add count from last row id to dense_shape - 1 new_sparse_indices_shape[0] = int64(count) new_sparse_indices_shape[1] = int64(sparse_indices.shape[1]) new_sparse_values_shape[0] = int64(count) From e2d3d45604d46e1aca76dc0f53cd42612ab2fc2d Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 15 Feb 2021 21:21:55 +0000 Subject: [PATCH 29/32] Documentation --- src/relay/op/tensor/transform.cc | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index ae6cebd99878..41bf77409475 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1587,7 +1587,8 @@ RELAY_REGISTER_OP("repeat") bool SparseFillEmptyRowsRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { // types: [sparse_indices, sparse_values, dense_shape, default_value, result] - ICHECK_EQ(types.size(), 5); + ICHECK_EQ(types.size(), 5) << "SparseFillEmptyRowsRel expects 5 inputs but " << types.size() + << "provided"; std::vector fields; auto sparse_indices = types[0].as(); auto ndims = sparse_indices->shape[1]; @@ -1613,14 +1614,17 @@ RELAY_REGISTER_OP("sparse_fill_empty_rows") .set_num_inputs(4) .add_argument("sparse_indices", "Tensor", "A 2-D int64 tensor of shape [N, ndims], which specifies the indices of the" - "elements in the sparse tensor that contain nonzero values") - .add_argument("sparse_values", "Tensor", - "A 1-D tensor[N] which supplies the values for each element in indices") + "elements in the sparse tensor that contain nonzero values. COO Format") + .add_argument( + "sparse_values", "Tensor", + "A 1-D tensor[N] which supplies the values for each element in indices. COO Format") .add_argument("dense_shape", "Tensor", "A 1-D int64 tensor of shape [ndims], which specifies the dense_shape of the" - "sparse tensor. Takes a list indicating the number of elements in each dimension") - .add_argument("default_value", "Tensor", - "The value to fill for empty rows, with the same type as sparse_values") + "sparse tensor. Takes a list indicating the number of elements in each " + "dimension. COO Format") + .add_argument( + "default_value", "Tensor", + "The value to fill for empty rows, with the same type as sparse_values. COO Format") .add_type_rel("sparse_fill_empty_rows", SparseFillEmptyRowsRel) .set_support_level(3) .set_attr("TOpPattern", kOpaque); From dcc2d384c83b2ae77d9b47b56b66b2e2f5b95184 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 16 Feb 2021 20:47:01 +0000 Subject: [PATCH 30/32] PR Changes --- python/tvm/relay/op/transform.py | 10 +++-- python/tvm/topi/sparse_fill_empty_rows.py | 47 ++++++++--------------- python/tvm/topi/transform.py | 10 +++-- src/relay/op/tensor/transform.cc | 7 ++-- 4 files changed, 31 insertions(+), 43 deletions(-) diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 245b4ec4e77d..ebad05e84f2c 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1324,14 +1324,15 @@ def adv_index(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. + Fill rows in a sparse matrix that do no contain any values. Values are placed in the first + column of empty rows. The sparse array is in COO format. It returns a TupleWrapper with 3 outputs Parameters ---------- sparse_indices : relay.Expr A 2-D int64 tensor[N, ndims] 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. - The inputs for this relay parameter must be in row-major order. + The first column of this relay parameter must be sorted in ascending order. sparse_values : relay.Expr A 1-D int64 tensor[N] containing the sparse values for the sparse indices. dense_shape : relay.Expr @@ -1342,14 +1343,15 @@ def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_v ------- new_sparse_indices : relay.Expr A 2-D int64 tensor[?, ndims] of integers containing location of new sparse - indices + indices. The first column outputs must be sorted in ascending order. new_sparse_values : relay.Expr A 1-D int64 tensor[?] containing the sparse values for the sparse indices. empty_row_indicator : relay.Expr A 1-D int64 tensor[dense_shape[0]] filled with zeros and ones indicating whether the particular row is empty or full respectively - Note: + Note + ---- This op exactly follows the documentation here: https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows There are two exceptions: diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index 8aa5f9b42314..3123b9ecd8d8 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -49,46 +49,31 @@ def _sparse_fill_empty_rows( return (new_sparse_indices, new_sparse_values, empty_row_indicator) else: - # Add rows with default value if first row id of sparse_indices is not a zero - for i in range(0, int64(sparse_indices[0, 0])): - new_sparse_indices[new_sparse_indices_row_id, 0] = int64(i) - for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[new_sparse_indices_row_id, k] = int64(0) - - new_sparse_values[new_sparse_indices_row_id] = default_value_ - empty_row_indicator[i] = int64(1) - new_sparse_indices_row_id += 1 - # Iterate through sparse_indices and add rows if/when required for i in range(0, int64(sparse_indices.shape[0])): - row_id = int64(sparse_indices[i, 0]) if i == 0: - # Add first row of input to output - new_sparse_indices[new_sparse_indices_row_id, 0] = row_id - for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[new_sparse_indices_row_id, k] = int64(sparse_indices[i, k]) - new_sparse_values[new_sparse_indices_row_id] = int64(sparse_values[i]) - empty_row_indicator[row_id] = int64(0) - new_sparse_indices_row_id += 1 + prev_row_id = int64(0) else: prev_row_id = int64(sparse_indices[i - 1, 0] + 1) - # Since input is in row-major order, add rows between prev_row_id and row_id - for j in range(prev_row_id, row_id): - new_sparse_indices[new_sparse_indices_row_id, 0] = int64(j) - for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[new_sparse_indices_row_id, k] = int64(0) - empty_row_indicator[prev_row_id] = int64(1) - new_sparse_values[new_sparse_indices_row_id] = default_value_ - new_sparse_indices_row_id += 1 + row_id = int64(sparse_indices[i, 0]) - # Add current row to output - new_sparse_indices[new_sparse_indices_row_id, 0] = row_id + # Since input is in row-major order, add rows between prev_row_id and row_id + for j in range(prev_row_id, row_id): + new_sparse_indices[new_sparse_indices_row_id, 0] = int64(j) for k in range(1, int64(new_sparse_indices_shape[1])): - new_sparse_indices[new_sparse_indices_row_id, k] = int64(sparse_indices[i, k]) - new_sparse_values[new_sparse_indices_row_id] = int64(sparse_values[i]) - empty_row_indicator[row_id] = int64(0) + new_sparse_indices[new_sparse_indices_row_id, k] = int64(0) + empty_row_indicator[prev_row_id] = int64(1) + new_sparse_values[new_sparse_indices_row_id] = default_value_ new_sparse_indices_row_id += 1 + # Add current element to output + new_sparse_indices[new_sparse_indices_row_id, 0] = row_id + for k in range(1, int64(new_sparse_indices_shape[1])): + new_sparse_indices[new_sparse_indices_row_id, k] = int64(sparse_indices[i, k]) + new_sparse_values[new_sparse_indices_row_id] = int64(sparse_values[i]) + empty_row_indicator[row_id] = int64(0) + new_sparse_indices_row_id += 1 + # Add rows with default value if last row id of sparse_indices is not dense_shape[0] - 1 for i in range( int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1), int64(dense_shape[0]) diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index 794fa53c786f..cc390657a5e6 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -935,14 +935,15 @@ def adv_index(data, indices): 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. + Fill rows in a sparse matrix that do no contain any values. Values are placed in the first + column of empty rows. The sparse array is in COO format. It returns a TupleWrapper with 3 outputs Parameters ---------- sparse_indices : relay.Expr A 2-D int64 tensor[N, ndims] 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. - The inputs for this relay parameter must be in row-major order. + The first column of this relay parameter must be sorted in ascending order. sparse_values : relay.Expr A 1-D int64 tensor[N] containing the sparse values for the sparse indices. dense_shape : relay.Expr @@ -953,14 +954,15 @@ def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_v ------- new_sparse_indices : relay.Expr A 2-D int64 tensor[?, ndims] of integers containing location of new sparse - indices + indices. The first column outputs must be sorted in ascending order. new_sparse_values : relay.Expr A 1-D int64 tensor[?] containing the sparse values for the sparse indices. empty_row_indicator : relay.Expr A 1-D int64 tensor[dense_shape[0]] filled with zeros and ones indicating whether the particular row is empty or full respectively - Note: + Note + ---- This op exactly follows the documentation here: https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows There are two exceptions: diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 41bf77409475..35954746ab1f 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1621,10 +1621,9 @@ RELAY_REGISTER_OP("sparse_fill_empty_rows") .add_argument("dense_shape", "Tensor", "A 1-D int64 tensor of shape [ndims], which specifies the dense_shape of the" "sparse tensor. Takes a list indicating the number of elements in each " - "dimension. COO Format") - .add_argument( - "default_value", "Tensor", - "The value to fill for empty rows, with the same type as sparse_values. COO Format") + "dimension") + .add_argument("default_value", "Tensor", + "The value to fill for empty rows, with the same type as sparse_values") .add_type_rel("sparse_fill_empty_rows", SparseFillEmptyRowsRel) .set_support_level(3) .set_attr("TOpPattern", kOpaque); From dd055b3632e06302eeceec8adb0aa97c7df21555 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 16 Feb 2021 21:00:14 +0000 Subject: [PATCH 31/32] PR Comments --- python/tvm/topi/sparse_fill_empty_rows.py | 4 ++-- tests/python/relay/dyn/test_dynamic_op_level3.py | 12 ++++++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index 3123b9ecd8d8..aad936411248 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -32,8 +32,8 @@ def _sparse_fill_empty_rows( empty_row_indicator_shape, ): default_value_ = int64(default_value[0]) - new_sparse_indices = output_tensor(new_sparse_indices_shape, sparse_indices.dtype) - new_sparse_values = output_tensor(new_sparse_values_shape, sparse_values.dtype) + new_sparse_indices = output_tensor(new_sparse_indices_shape, "int64") + new_sparse_values = output_tensor(new_sparse_values_shape, "int64") empty_row_indicator = output_tensor(empty_row_indicator_shape, "int64") new_sparse_indices_row_id = 0 diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 12b15f06e976..689dfd4729a0 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -244,8 +244,11 @@ def verify_sparse_to_dense(sparse_indices, sparse_values, default_value, output_ ), ], ) +@pytest.mark.parametrize("dtype", [np.int32, np.int64]) @pytest.mark.parametrize("use_dyn", [True, False]) -def test_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value, use_dyn): +def test_sparse_fill_empty_rows( + sparse_indices, sparse_values, dense_shape, default_value, dtype, use_dyn +): def ref_sparse_fill_empty_rows( sparse_indices: np.ndarray, sparse_values: np.ndarray, @@ -345,7 +348,12 @@ def verify_sparse_fill_empty_rows( [("llvm", tvm.cpu())], ) - verify_sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) + verify_sparse_fill_empty_rows( + sparse_indices.astype(dtype), + sparse_values.astype(dtype), + dense_shape.astype(dtype), + default_value.astype(dtype), + ) if __name__ == "__main__": From c43cda2d474d7f72a1274e688f1154095bcdd787 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 16 Feb 2021 23:15:35 +0000 Subject: [PATCH 32/32] Resolve input and output dtype compat --- python/tvm/relay/frontend/tensorflow.py | 4 +- python/tvm/relay/op/transform.py | 21 +++--- python/tvm/topi/sparse_fill_empty_rows.py | 2 +- python/tvm/topi/transform.py | 65 ------------------- .../relay/dyn/test_dynamic_op_level3.py | 20 ++++-- 5 files changed, 29 insertions(+), 83 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 6d6d452d662c..8529205d96d4 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -1013,9 +1013,7 @@ def _impl(inputs, attr, params, mod): ) return _expr.TupleWrapper( - _expr.Tuple( - [new_sparse_indices, new_sparse_values, _op.cast(empty_row_indicator, dtype="bool")] - ), + _expr.Tuple([new_sparse_indices, new_sparse_values, empty_row_indicator]), 3, ) diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index ebad05e84f2c..f8eed3516f91 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1330,24 +1330,24 @@ def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_v Parameters ---------- sparse_indices : relay.Expr - A 2-D int64 tensor[N, ndims] of integers containing location of sparse values, where N is + A 2-D tensor[N, ndims] 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. The first column of this relay parameter must be sorted in ascending order. sparse_values : relay.Expr - A 1-D int64 tensor[N] containing the sparse values for the sparse indices. + A 1-D tensor[N] containing the sparse values for the sparse indices. dense_shape : relay.Expr - A 1-D int64 tensor[ndims] which contains shape of the dense output tensor. + A 1-D tensor[ndims] which contains shape of the dense output tensor. default_value : relay.Expr A 1-D tensor[1] containing the default value for the remaining locations. Returns ------- new_sparse_indices : relay.Expr - A 2-D int64 tensor[?, ndims] of integers containing location of new sparse + A 2-D tensor[?, ndims] of integers containing location of new sparse indices. The first column outputs must be sorted in ascending order. new_sparse_values : relay.Expr - A 1-D int64 tensor[?] containing the sparse values for the sparse indices. + A 1-D tensor[?] containing the sparse values for the sparse indices. empty_row_indicator : relay.Expr - A 1-D int64 tensor[dense_shape[0]] filled with zeros and ones + A 1-D tensor[dense_shape[0]] filled with zeros and ones indicating whether the particular row is empty or full respectively Note @@ -1380,13 +1380,18 @@ def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_v [2, 0], [3, 1], [4, 0]] - empty_row_indicator = [0, 1, 0, 0, 1] + empty_row_indicator = [False, True, False, False, True] new_sparse_values = [1, 2, 10, 3, 4, 10] """ - return TupleWrapper( + new_sparse_indices, new_sparse_values, empty_row_indicator = TupleWrapper( _make.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value), 3 ) + new_sparse_indices = cast_like(new_sparse_indices, sparse_indices) + new_sparse_values = cast_like(new_sparse_values, sparse_values) + empty_row_indicator = cast(empty_row_indicator, "bool") + + return Tuple((new_sparse_indices, new_sparse_values, empty_row_indicator)) def cumsum(data, axis=None, dtype=None, exclusive=None): diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py index aad936411248..10dc6ee3bfa3 100644 --- a/python/tvm/topi/sparse_fill_empty_rows.py +++ b/python/tvm/topi/sparse_fill_empty_rows.py @@ -39,7 +39,7 @@ def _sparse_fill_empty_rows( if int64(sparse_indices.shape[0]) == int64(0): # Handle Empty Case # Fill all rows with default values - for i in range(0, int64(new_sparse_indices_shape[0])): + for i in range(0, new_sparse_indices_shape[0]): new_sparse_indices[i, 0] = int64(i) new_sparse_values[i] = default_value_ empty_row_indicator[i] = int64(1) diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index cc390657a5e6..6ddbc73e4666 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -931,68 +931,3 @@ def adv_index(data, indices): Output tensor """ return cpp.adv_index(data, indices) - - -def sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value): - """ - Fill rows in a sparse matrix that do no contain any values. Values are placed in the first - column of empty rows. The sparse array is in COO format. - It returns a TupleWrapper with 3 outputs - Parameters - ---------- - sparse_indices : relay.Expr - A 2-D int64 tensor[N, ndims] 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. - The first column of this relay parameter must be sorted in ascending order. - sparse_values : relay.Expr - A 1-D int64 tensor[N] containing the sparse values for the sparse indices. - dense_shape : relay.Expr - A 1-D int64 tensor[ndims] which contains shape of the dense output tensor. - default_value : relay.Expr - A 1-D tensor[1] containing the default value for the remaining locations. - Returns - ------- - new_sparse_indices : relay.Expr - A 2-D int64 tensor[?, ndims] of integers containing location of new sparse - indices. The first column outputs must be sorted in ascending order. - new_sparse_values : relay.Expr - A 1-D int64 tensor[?] containing the sparse values for the sparse indices. - empty_row_indicator : relay.Expr - A 1-D int64 tensor[dense_shape[0]] filled with zeros and ones - indicating whether the particular row is empty or full respectively - - Note - ---- - This op exactly follows the documentation here: - https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows - There are two exceptions: - 1. Input Sparse Indices are expected to be in row-major order. - 2. Empty Row Indicator has int64 output type with 1(for True) and 0(for False). - - Examples - ------- - .. code-block:: python - sparse_indices = [[0, 1], - [0, 3], - [2, 0], - [3, 1]] - sparse_values = [1, 2, 3, 4] - default_value = [10] - dense_shape = [5, 6] - new_sparse_indices, empty_row_indicator, new_sparse_values, slice_element_index = - relay.sparse_fill_empty_rows( - sparse_indices, - sparse_values, - default_value, - dense_shape) - new_sparse_indices = [[0, 1], - [0, 3], - [1, 0], - [2, 0], - [3, 1], - [4, 0]] - empty_row_indicator = [0, 1, 0, 0, 1] - new_sparse_values = [1, 2, 10, 3, 4, 10] - - """ - return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 689dfd4729a0..d5f81e84e39d 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -244,7 +244,7 @@ def verify_sparse_to_dense(sparse_indices, sparse_values, default_value, output_ ), ], ) -@pytest.mark.parametrize("dtype", [np.int32, np.int64]) +@pytest.mark.parametrize("dtype", [np.int64, np.int32]) @pytest.mark.parametrize("use_dyn", [True, False]) def test_sparse_fill_empty_rows( sparse_indices, sparse_values, dense_shape, default_value, dtype, use_dyn @@ -264,7 +264,7 @@ def check_add_rows(current_idx, limit_idx): while current_idx < limit_idx: new_sparse_indices.append([current_idx] + [0] * (num_cols - 1)) new_sparse_values.append(default_value[0]) - empty_row_indicator[current_idx] = 1 + empty_row_indicator[current_idx] = True current_idx += 1 return current_idx @@ -272,7 +272,7 @@ def check_add_rows(current_idx, limit_idx): current_idx = 0 new_sparse_indices = [] new_sparse_values = [] - empty_row_indicator = [0 for _ in range(dense_shape[0])] + empty_row_indicator = [False for _ in range(dense_shape[0])] num_cols = sparse_indices.shape[1] for sparse_row, sparse_value in zip(sparse_indices, sparse_values): limit_idx = sparse_row[0] @@ -331,9 +331,7 @@ def verify_sparse_fill_empty_rows( "default_value", relay.TensorType(default_value_np.shape, str(default_value_np.dtype)), ) - z = relay.sparse_fill_empty_rows( - sparse_indices, sparse_values, dense_shape, default_value - ).astuple() + z = relay.sparse_fill_empty_rows(sparse_indices, sparse_values, dense_shape, default_value) func = relay.Function([sparse_indices, sparse_values, dense_shape, default_value], z) ref_res = ref_sparse_fill_empty_rows( sparse_indices_np, @@ -341,6 +339,16 @@ def verify_sparse_fill_empty_rows( dense_shape_np, default_value_np, ) + ( + new_sparse_indices_infer_type, + new_sparse_values_infer_type, + empty_row_indicator_infer_type, + ) = run_infer_type(z) + + assert new_sparse_indices_infer_type.checked_type.dtype == sparse_indices_np.dtype + assert new_sparse_values_infer_type.checked_type.dtype == sparse_indices_np.dtype + assert empty_row_indicator_infer_type.checked_type.dtype == "bool" + verify_func( func, [sparse_indices_np, sparse_values_np, dense_shape_np, default_value_np],