diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index cbe989f93558..68688dd61064 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -401,6 +401,14 @@ struct SparseToDenseAttrs : public tvm::AttrsNode { } }; // struct SparseToDenseAttrs +/*! \brief Attributes used in sparsefillemptyRows operator */ +struct SparseFillEmptyRowsAttrs : public tvm::AttrsNode { + Array dense_shape; + + TVM_DECLARE_ATTRS(SparseFillEmptyRowsAttrs, "relay.attrs.SparseFillEmptyRowsAttrs") { + TVM_ATTR_FIELD(dense_shape).describe("Shape of the dense output tensor"); + } +}; // struct SparseFillEmptyRowsAttrs /*! \brief Attributes for ndarray_size operator */ struct NdarraySizeAttrs : public tvm::AttrsNode { DataType dtype; diff --git a/include/tvm/support/logging.h b/include/tvm/support/logging.h index d98363ea1c1b..ced1902a1bd1 100644 --- a/include/tvm/support/logging.h +++ b/include/tvm/support/logging.h @@ -139,10 +139,10 @@ constexpr const char* kTVM_INTERNAL_ERROR_MESSAGE = #define ICHECK_GE(x, y) ICHECK_BINARY_OP(_GE, >=, x, y) #define ICHECK_EQ(x, y) ICHECK_BINARY_OP(_EQ, ==, x, y) #define ICHECK_NE(x, y) ICHECK_BINARY_OP(_NE, !=, x, y) -#define ICHECK_NOTNULL(x) \ - ((x) == nullptr ? dmlc::LogMessageFatal(__FILE__, __LINE__).stream() \ - << tvm::kTVM_INTERNAL_ERROR_MESSAGE << __INDENT << "Check not null: " #x \ - << ' ', \ +#define ICHECK_NOTNULL(x) \ + ((x) == nullptr ? dmlc::LogMessageFatal(__FILE__, __LINE__).stream() \ + << tvm::kTVM_INTERNAL_ERROR_MESSAGE << ICHECK_INDENT \ + << "Check not null: " #x << ' ', \ (x) : (x)) // NOLINT(*) /*! \brief The diagnostic level, controls the printing of the message. */ diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index a04762f28feb..46029299c2eb 100644 --- a/include/tvm/topi/transform.h +++ b/include/tvm/topi/transform.h @@ -1386,6 +1386,94 @@ inline Array meshgrid(const Array& inputs, const std::string& in return result; } +/*! + * \brief Fill empty rows of a sparse tensor with default values + * + * \param sparse_indices Indices where values of the dense tensor exist + * \param sparse_values Values at the above indices respectively + * \param default_value Default value at to be used at empty rows + * \param dense_shape Dense shape of the sparse tensor + * \param name The name of the operation + * \param tag The tag to mark the operation + * + * \return A Tensor whose op member is the SparseFillEmptyRows operation + */ +inline Array SparseFillEmptyRows(const Tensor& sparse_indices, const Tensor& sparse_values, + const Tensor& default_value, + const Array& dense_shape, + const std::string name = "T_sparse_fill_empty_rows", + std::string tag = kInjective) { + Array result; + Array sp_ordered_output_shape; + sp_ordered_output_shape.push_back(dense_shape[0] + sparse_indices->shape[0]); + if (sparse_indices->shape.size() > 1) { + sp_ordered_output_shape.push_back(sparse_indices->shape[1]); + } + auto empty_row_indicator = + compute(Array{dense_shape[0]}, [&](const Array& indices) { + PrimExpr ret = PrimExpr(Bool(1)); + for (int i = 0; i < GetConstInt(sparse_indices->shape[0]); ++i) { + PrimExpr sparse_index; + if (sparse_indices->shape.size() == 1) { + sparse_index = sparse_indices[i]; + } else { + sparse_index = sparse_indices[i][0]; + } + ret = if_then_else(sparse_index == indices[0], PrimExpr(Bool(0)), ret); + } + return ret; + }); + result.push_back(compute( + sp_ordered_output_shape, + [&](const Array& indices) { + PrimExpr ret = -1; + ret = if_then_else(indices[0] < sparse_indices->shape[0], sparse_indices(indices), ret); + PrimExpr empty_row_count = 0; + for (int i = 0; i < static_cast(dense_shape[0]); ++i) { + empty_row_count = + if_then_else(empty_row_indicator[i], empty_row_count + 1, empty_row_count); + PrimExpr at_correct_index = + (indices[0] == (sparse_indices->shape[0] + empty_row_count - 1)); + PrimExpr condition = at_correct_index && empty_row_indicator[i]; + + ret = if_then_else(condition, i, ret); + if (indices.size() > 1) { + ret = if_then_else(condition && indices[1] > 0, 0, ret); + } + } + return ret; + }, + name, tag)); + result.push_back(empty_row_indicator); + result.push_back(compute( + Array{sp_ordered_output_shape[0]}, + [&](const Array& indices) { + PrimExpr ret = -1; + ret = if_then_else(indices[0] < sparse_values->shape[0], sparse_values(indices), ret); + PrimExpr empty_row_count = 0; + for (int i = 0; i < static_cast(dense_shape[0]); ++i) { + empty_row_count = + if_then_else(empty_row_indicator[i], empty_row_count + 1, empty_row_count); + PrimExpr condition = + (indices[0] == sparse_values->shape[0] + empty_row_count - 1) && empty_row_count > 0; + ret = if_then_else(condition, default_value[0], ret); + } + return ret; + }, + name, tag)); + result.push_back(compute( + Array{1}, + [&](const Array& indices) { + PrimExpr non_empty_rows = 0; + for (int i = 0; i < static_cast(dense_shape[0]); ++i) { + non_empty_rows = if_then_else(empty_row_indicator[i], non_empty_rows, non_empty_rows + 1); + } + return non_empty_rows; + }, + name, tag)); + return result; +} + /*! * \brief Transform the layout according to \p src_layout and \p dst_layout * \param src the source input. diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index d5746a38582c..5b6f64bf58b6 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -1365,6 +1365,54 @@ 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" + + indices_tensor = _infer_value(inputs[0], params, mod).asnumpy() + values_tensor = _infer_value(inputs[1], params, mod).asnumpy() + dense_shape_tensor = _infer_value(inputs[2], params, mod).asnumpy() + default_value_tensor = _infer_value(inputs[3], params, mod).asnumpy().reshape(1) + + indices_data = _expr.const(indices_tensor, indices_tensor.dtype) + values_data = _expr.const(values_tensor, values_tensor.dtype) + default_value_data = _expr.const(default_value_tensor, default_value_tensor.dtype) + + ( + new_sparse_indices, + empty_row_indicator, + new_sparse_values, + non_empty_rows, + ) = get_relay_op("sparse_fill_empty_rows")( + indices_data, values_data, default_value_data, list(dense_shape_tensor) + ) + first_column = get_relay_op("split")(new_sparse_indices, indices_tensor.shape[1], axis=1) + sorted_indices = _op.argsort(_op.squeeze(first_column[0])) + + final_sparse_indices = _op.strided_slice( + _op.take(new_sparse_indices, sorted_indices, axis=0), + begin=_op.concatenate([non_empty_rows, _expr.const([0])], 0), + end=[-1, -1], + strides=[1, 1], + slice_mode="size", + ) + + final_sparse_values = _op.strided_slice( + _op.take(new_sparse_values, sorted_indices), + begin=non_empty_rows, + end=_expr.const([-1]), + slice_mode="size", + ) + + return ( + final_sparse_indices, + final_sparse_values, + empty_row_indicator, + ) + + return _impl + + def _bias_add(): def _impl(inputs, attr, params, mod): # Must expand for proper broadcasting in NCHW. @@ -2422,6 +2470,7 @@ def _impl(inputs, attr, params, mod): "SpaceToBatchND": _space_to_batch_nd(), "SpaceToDepth": _space_to_depth(), "SparseToDense": _sparse_to_dense(), + "SparseFillEmptyRows": _sparse_fill_empty_rows(), "SparseTensorDenseMatMul": _sparse_tensor_dense_matmul(), "Split": _split(False), "SplitV": _split(True), diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 05ca6d2e4bb9..f66d5ea76c4f 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -63,6 +63,7 @@ _reg.register_injective_schedule("sparse_to_dense") _reg.register_injective_schedule("matrix_set_diag") _reg.register_injective_schedule("adv_index") +_reg.register_injective_schedule("sparse_fill_empty_rows") # concatenate _reg.register_schedule("concatenate", strategy.schedule_concatenate) diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 7e7f9b299593..e2314e6f81f1 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1320,3 +1320,83 @@ def adv_index(inputs): Output tensor. """ return _make.adv_index(Tuple(inputs)) + + +def sparse_fill_empty_rows(sparse_indices, sparse_values, default_value, dense_shape): + """ + Fill first column of the empty rows with default values for a sparse array. + It returns a TupleWrapper with four outputs + + Parameters + ---------- + sparse_indices : relay.Expr + A 2-D tensor[N, n_dim] of integers containing location of sparse values, where N is the + number of sparse values and n_dim is the number of dimensions of the dense_shape + + sparse_values : relay.Expr + A 1-D tensor[N] containing the sparse values for the sparse indices. + + default_value : relay.Expr + A 1-D tensor containing the default value for the remaining locations. + + dense_shape : relay.Expr + A list of integers. Shape of the dense output tensor. + + Returns + ------- + new_sparse_indices : relay.Expr + A 2-D tensor[N + dense_shape[0], n_dim] of integers containing location of new sparse + indices where N is the number of sparse values. It is filled with -1 at irrelevant indices + which will be sliced in a future op discarding non-useful elements. This is done since the + real rows of new_sparse_indices depends on the input. + + empty_row_indicator : relay.Expr + A 1-D Boolean tensor[dense_shape[0]] indicating whether the particular row is empty + + new_sparse_values : relay.Expr + A 1-D tensor[dense_shape[0]] containing the sparse values for the sparse indices. It is + filled with -1 at indices which will be discarded in the following strided_slice op. + This is done since the real rows of new_sparse_indices depends on the input. + + non_empty_rows : relay.Expr + A 1-D tensor containing the amount of non-empty rows in the sparse_indices. This value will + be used to slice irrelevant indices(filled with -1) in new_sparse_values and + new_sparse_indices + + 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], + [2, 0], + [3, 1], + [1, 0], + [4, 0], + [-1, -1], + [-1, -1], + [-1, -1]] + + empty_row_indicator = [False, True, False, False, True] + + new_sparse_values = [1, 2, 3, 4, 10, 10, -1, -1, -1] + + slice_element_index = [3] + + """ + return TupleWrapper( + _make.sparse_fill_empty_rows(sparse_indices, sparse_values, default_value, dense_shape), 4 + ) diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index 6ddbc73e4666..c759c3d09960 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -931,3 +931,81 @@ def adv_index(data, indices): Output tensor """ return cpp.adv_index(data, indices) + + +def sparse_fill_empty_rows(sparse_indices, sparse_values, default_value, dense_shape): + """ + Fill first column of the empty rows with default values for a sparse array. + It returns a TupleWrapper with four outputs + + Parameters + ---------- + sparse_indices : relay.Expr + A 2-D tensor[N, n_dim] of integers containing location of sparse values, where N is the + number of sparse values and n_dim is the number of dimensions of the dense_shape + + sparse_values : relay.Expr + A 1-D tensor[N] containing the sparse values for the sparse indices. + + default_value : relay.Expr + A 1-D tensor containing the default value for the remaining locations. + + dense_shape : relay.Expr + A list of integers. Shape of the dense output tensor. + + Returns + ------- + new_sparse_indices : relay.Expr + A 2-D tensor[N + dense_shape[0], n_dim] of integers containing location of new sparse + indices where N is the number of sparse values. It is filled with -1 at irrelevant indices + which will be sliced in a future op discarding non-useful elements. This is done since the + real rows of new_sparse_indices depends on the input. + + empty_row_indicator : relay.Expr + A 1-D Boolean tensor[dense_shape[0]] indicating whether the particular row is empty + + new_sparse_values : relay.Expr + A 1-D tensor[dense_shape[0]] containing the sparse values for the sparse indices. It is + filled with -1 at indices which will be discarded in the following strided_slice op. + This is done since the real rows of new_sparse_indices depends on the input. + + non_empty_rows : relay.Expr + A 1-D tensor containing the amount of non-empty rows in the sparse_indices. This value will + be used to slice irrelevant indices(filled with -1) in new_sparse_values and + new_sparse_indices + + 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], + [2, 0], + [3, 1], + [1, 0], + [4, 0], + [-1, -1], + [-1, -1], + [-1, -1]] + + empty_row_indicator = [False, True, False, False, True] + + new_sparse_values = [1, 2, 3, 4, 10, 10, -1, -1, -1] + + slice_element_index = [3] + + """ + return cpp.sparse_fill_empty_rows(sparse_indices, sparse_values, default_value, dense_shape) diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 640943eac805..49bfefb2bf6f 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -1553,6 +1553,70 @@ RELAY_REGISTER_OP("meshgrid") .set_attr("FTVMCompute", MeshgridCompute) .set_attr("TOpPattern", kInjective); +TVM_REGISTER_NODE_TYPE(SparseFillEmptyRowsAttrs); + +bool SparseFillEmptyRowsRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + // types: [ sparse_indices, sparse_values, default_value, result] + ICHECK_EQ(types.size(), 4) << "SparseFillEmptyRowsRel expects 4 arguments but " << types.size() + << " were provided."; + std::vector fields; + auto sparse_indices = types[0].as(); + auto default_value = types[2].as(); + const auto* param = attrs.as(); + ICHECK_NOTNULL(param); + + Array sp_ordered_output_shape; + sp_ordered_output_shape.push_back(param->dense_shape[0] + sparse_indices->shape[0]); + if (sparse_indices->shape.size() > 1) { + sp_ordered_output_shape.push_back(sparse_indices->shape[1]); + } + fields.push_back(TensorType(sp_ordered_output_shape, sparse_indices->dtype)); + fields.push_back(TensorType(Array{param->dense_shape[0]}, tvm::DataType::Bool())); + fields.push_back(TensorType(Array{sp_ordered_output_shape[0]}, default_value->dtype)); + fields.push_back(TensorType(Array{1}, tvm::DataType::Int(32))); + reporter->Assign(types[3], TupleType(Array(fields))); + return true; +} + +Array SparseFillEmptyRowsCompute(const Attrs& attrs, const Array& inputs, + const Type& out_type) { + ICHECK_EQ(inputs.size(), 3) << "SparseFillEmptyRowsCompute expects 3 arguments but " + << inputs.size() << " were provided."; + const auto* param = attrs.as(); + ICHECK_NOTNULL(param); + return {topi::SparseFillEmptyRows(inputs[0], inputs[1], inputs[2], param->dense_shape)}; +} + +Expr MakeSparseFillEmptyRows(Expr sparse_indices, Expr sparse_values, Expr default_value, + Array dense_shape) { + auto attrs = make_object(); + attrs->dense_shape = std::move(dense_shape); + static const Op& op = Op::Get("sparse_fill_empty_rows"); + return Call(op, {sparse_indices, sparse_values, default_value}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.sparse_fill_empty_rows") + .set_body_typed(MakeSparseFillEmptyRows); + +RELAY_REGISTER_OP("sparse_fill_empty_rows") + .describe(R"code(Return representation of a sparse tensor with empty rows filled with default + value.)code" TVM_ADD_FILELINE) + .set_num_inputs(3) + .set_attrs_type() + .add_argument( + "sparse_indices", "Tensor", + "A 2-D tensor[N, n_dim] of integers containing location of sparse values, where N is the" + "number of sparse values and n_dim is the number of dimensions of the dense_shape") + .add_argument("sparse_values", "Tensor", + "A 1-D tensor[N] containing the sparse values for the sparse indices") + .add_argument("default_value", "Tensor", + "A tensor containing the default value for the remaining locations") + .add_type_rel("sparse_fill_empty_rows", SparseFillEmptyRowsRel) + .set_support_level(3) + .set_attr("TOpPattern", kInjective) + .set_attr("FTVMCompute", SparseFillEmptyRowsCompute); + // tile operator TVM_REGISTER_NODE_TYPE(TileAttrs); diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 22ed6c5b2edf..40048b52bca9 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -1811,6 +1811,73 @@ def test_forward_sparse_dense_matmul(): ) +####################################################################### +# SparseFillEmptyRows +# ------------ + + +def _test_sparse_fill_empty_rows(indices_np, values_np, default_value, dense_shape_np): + with tf.Graph().as_default(): + sp_input = tf.sparse.SparseTensor( + indices=indices_np, values=values_np, dense_shape=dense_shape_np + ) + result = tf.sparse.fill_empty_rows(sp_input, default_value, name="sparse_fill_empty_rows") + compare_tf_with_tvm( + None, + "", + [ + "sparse_fill_empty_rows/SparseFillEmptyRows:0", + "sparse_fill_empty_rows/SparseFillEmptyRows:1", + "sparse_fill_empty_rows/SparseFillEmptyRows:2", + "SparseTensor/dense_shape:0", + ], + ) + + +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.int32) + sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int32) + dense_shape_np = np.array([5, 6], dtype=np.int32) + default_value = 10 + _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, default_value, dense_shape_np) + + sparse_indices_np = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int32) + sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int32) + dense_shape_np = np.array([7, 7, 7], dtype=np.int32) + default_value_np = np.array([10], dtype=np.int32) + _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, default_value, dense_shape_np) + + sparse_indices_np = np.array([[1], [2]], dtype=np.int32) + sparse_values_np = np.array([7, 8], dtype=np.int32) + dense_shape_np = np.array([5], dtype=np.int32) + default_value_np = np.array([4], dtype=np.int32) + _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, default_value, dense_shape_np) + + sparse_indices_np = np.ones((0, 1), dtype=np.int32) + sparse_values_np = np.array([], dtype=np.int32) + dense_shape_np = np.array([5], dtype=np.int32) + default_value_np = np.array([4], dtype=np.int32) + _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, default_value, dense_shape_np) + + sparse_indices_np = np.ones((0, 3), dtype=np.int32) + sparse_values_np = np.array([], dtype=np.int32) + dense_shape_np = np.array([9, 3, 7], dtype=np.int32) + default_value_np = np.array([100], dtype=np.int32) + _test_sparse_fill_empty_rows(sparse_indices_np, sparse_values_np, default_value, dense_shape_np) + + ####################################################################### # StridedSlice # ------------ diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index 668285dfb882..49ad8729be8d 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -1042,6 +1042,118 @@ def verify_scatter_add(dshape, ishape, axis=0, dtype="float32"): verify_scatter_add((16, 16, 4, 5), (16, 16, 4, 5), 3) +@tvm.testing.uses_gpu +def test_sparse_fill_empty_rows(): + def ref_sparse_fill_empty_rows( + sparse_indices: np.ndarray, + sparse_values: np.ndarray, + default_value: np.ndarray, + dense_shape: 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]) + slice_element_index = np.array(0, dtype=np.int32) + + 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 + else: + slice_element_index += 1 + + return new_sparse_indices, empty_row_indicator, new_sparse_values, slice_element_index + + def verify_sparse_fill_empty_rows( + sparse_indices_np: np.ndarray, + sparse_values_np: np.ndarray, + default_value_np: np.ndarray, + dense_shape_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)) + ) + default_value = relay.var( + "default_value", relay.TensorType(default_value_np.shape, str(default_value_np.dtype)) + ) + z = relay.op.sparse_fill_empty_rows( + sparse_indices, sparse_values, default_value, list(dense_shape_np) + ).astuple() + func = relay.Function([sparse_indices, sparse_values, default_value], z) + + ref_res = ref_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, default_value_np, dense_shape_np + ) + for target, ctx in tvm.testing.enabled_targets(): + for kind in ["graph", "debug"]: + intrp = relay.create_executor(kind, ctx=ctx, target=target) + op_res = intrp.evaluate(func)(sparse_indices_np, sparse_values_np, default_value_np) + for op_res_item, ref_res_item in zip(op_res, ref_res): + tvm.testing.assert_allclose(op_res_item.asnumpy(), ref_res_item, rtol=1e-5) + + sparse_indices_np = np.array([[0, 1], [0, 3], [2, 0], [3, 1]], dtype=np.int32) + sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int32) + dense_shape_np = np.array([5, 6], dtype=np.int32) + default_value_np = np.array([10], dtype=np.int32) + verify_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, default_value_np, dense_shape_np + ) + + sparse_indices_np = np.array([[1, 1, 1], [1, 3, 1], [2, 0, 5], [3, 1, 6]], dtype=np.int32) + sparse_values_np = np.array([1, 2, 3, 4], dtype=np.int32) + dense_shape_np = np.array([7, 7, 7], dtype=np.int32) + default_value_np = np.array([10], dtype=np.int32) + verify_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, default_value_np, dense_shape_np + ) + + sparse_indices_np = np.array([[1], [2]], dtype=np.int32) + sparse_values_np = np.array([7, 8], dtype=np.int32) + dense_shape_np = np.array([5], dtype=np.int32) + default_value_np = np.array([4], dtype=np.int32) + verify_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, default_value_np, dense_shape_np + ) + + sparse_indices_np = np.ones((0, 1), dtype=np.int32) + sparse_values_np = np.array([], dtype=np.int32) + dense_shape_np = np.array([5], dtype=np.int32) + default_value_np = np.array([4], dtype=np.int32) + verify_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, default_value_np, dense_shape_np + ) + + sparse_indices_np = np.ones((0, 3), dtype=np.int32) + sparse_values_np = np.array([], dtype=np.int32) + dense_shape_np = np.array([9, 3, 7], dtype=np.int32) + default_value_np = np.array([100], dtype=np.int32) + verify_sparse_fill_empty_rows( + sparse_indices_np, sparse_values_np, default_value_np, dense_shape_np + ) + + @tvm.testing.uses_gpu def test_gather(): def verify_gather(data, axis, indices, ref_res): @@ -1313,6 +1425,7 @@ def verify_adv_index(data_shape, index_shapes): if __name__ == "__main__": + test_sparse_fill_empty_rows() test_cast() test_zeros_ones() test_unary_identity()