From dd0ba5ac929d4b3c590125b87e34971ab4afdc47 Mon Sep 17 00:00:00 2001 From: Michalis Papadimitriou Date: Thu, 16 Jun 2022 14:14:56 +0300 Subject: [PATCH] [TOPI] Fix for trilu and set_matrix_diag --- include/tvm/relay/attrs/transform.h | 4 - include/tvm/topi/transform.h | 28 +++--- python/tvm/relay/frontend/onnx.py | 76 ++++++++++++---- python/tvm/relay/frontend/tflite.py | 18 ++-- python/tvm/relay/op/transform.py | 10 ++- python/tvm/topi/transform.py | 27 +++++- src/relay/op/tensor/transform.cc | 42 ++++----- src/topi/transform.cc | 5 +- tests/python/frontend/onnx/test_forward.py | 33 ++----- tests/python/relay/test_op_level10.py | 8 +- .../python/topi/python/test_topi_transform.py | 89 ++++++++++++++++++- 11 files changed, 242 insertions(+), 98 deletions(-) diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index 04c48a19ef22..9b3bf441e8b4 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -482,14 +482,10 @@ struct OneHotAttrs : public tvm::AttrsNode { /*! \brief Attributes used in matrix_set_diag operator */ struct MatrixSetDiagAttrs : public tvm::AttrsNode { - int k1; - int k2; bool super_diag_right_align; bool sub_diag_right_align; TVM_DECLARE_ATTRS(MatrixSetDiagAttrs, "relay.attrs.MatrixSetDiagAttrs") { - TVM_ATTR_FIELD(k1).set_default(0).describe("Lower limit (included) of the range of diagonals."); - TVM_ATTR_FIELD(k2).set_default(0).describe("Upper limit (included) of the range of diagonals."); TVM_ATTR_FIELD(super_diag_right_align) .set_default(true) .describe("Bool, true iff super-diagonal is right aligned (left-padded)."); diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index d7a1ef82f31a..c457a71977aa 100644 --- a/include/tvm/topi/transform.h +++ b/include/tvm/topi/transform.h @@ -244,7 +244,7 @@ inline Tensor transpose(const Tensor& x, Array axes, std::string name = * \param x The input tensor * \param seq_lengths A 1D Tensor with length x.dims[batch_axis]. Optional Tensor() can be passed. * If not defined batch axis is ignored and tensor is reversed along seq_axis. - * \param seq_axis The axis along which the elements will be reveresed + * \param seq_axis The axis along which the elements will be reversed * \param batch_axis The axis along which the tensor will be sliced * \param name The name of the operation * \param tag The tag to mark the operation @@ -267,7 +267,7 @@ inline Tensor reverse_sequence(const Tensor& x, const Tensor& seq_lengths, int s ICHECK(seq_lengths_dim == 1) << "seq_lengths should be 1D vector"; ICHECK(GetConstInt(seq_lengths->shape[0]) == GetConstInt(x->shape[batch_axis])) - << "For reverse_sequnece seq_lengths size should match with dimension of batch axis" + << "For reverse_sequence seq_lengths size should match with dimension of batch axis" << ", but got dimension of batch_axis = " << GetConstInt(x->shape[batch_axis]) << ", and seq_length size = " << GetConstInt(seq_lengths->shape[0]); @@ -763,7 +763,7 @@ inline Array StridedSliceOutputShape( * \param name The name of the operation * \param tag The tag to mark the operation * - * \return A Tensor whose op member is the sstrided_slice operation + * \return A Tensor whose op member is the strided_slice operation */ inline Tensor strided_slice_with_axes(const Tensor& x, const Array& begin, const Array& end, const Array& strides, @@ -1744,7 +1744,7 @@ inline Tensor ndarray_size(const Tensor& src, const DataType& dtype, } /*! - * \brief Returns a one-hot tensor where the locations repsented by indices take value on_value, + * \brief Returns a one-hot tensor where the locations represented by indices take value on_value, other locations take value off_value. * \param indices locations to set to on_value. * \param on_value value that locations represented by indices take on. @@ -1855,14 +1855,18 @@ inline Tensor sparse_to_dense(const Tensor& sparse_indices, const Arrayshape.size() - 1; - bool only_one_diagonal = k1 == k2; + std::cout << "\n input " << input->GetShape() << "\n" + << "diagonal " << diagonal << "\n k1 " << k1 << " \n k2 " << k2 + << "\n bool : " << only_one_diagonal; + return compute( input->shape, [&](const Array& iter_vars) { @@ -1873,11 +1877,12 @@ inline Tensor matrix_set_diag(const Tensor& input, const Tensor& diagonal, int k diagonal_indices.push_back(iter_vars[i]); } if (only_one_diagonal) { - k = k1; + k = k1(0); } else { // Determining which diagonal/sub-diagonal/super-diagonal it is k = iter_vars[ndim] - iter_vars[ndim - 1]; - diagonal_indices.push_back(k2 - k); + auto idx = k2(0) - k; + diagonal_indices.push_back(idx); // Calculating the offset in diagonal tensor for this diagonal auto get_offset = [&](PrimExpr M, PrimExpr N) { @@ -1895,8 +1900,9 @@ inline Tensor matrix_set_diag(const Tensor& input, const Tensor& diagonal, int k offset); return diagonal(diagonal_indices); }; - return if_then_else((PrimExpr)iter_vars[ndim] - iter_vars[ndim - 1] >= k1, - if_then_else((PrimExpr)iter_vars[ndim] - iter_vars[ndim - 1] <= k2, + + return if_then_else((PrimExpr)iter_vars[ndim] - iter_vars[ndim - 1] >= k1(0), + if_then_else((PrimExpr)iter_vars[ndim] - iter_vars[ndim - 1] <= k2(0), get_diag(), input(iter_vars)), input(iter_vars)); }, diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 595f12d4d5bd..7d17d817f28c 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -3487,7 +3487,8 @@ def get_var(name, val, scan=False): loop_vars = [ _expr.var(body.input[0].name, shape=(), dtype=iter_dtype), # iteration count - _expr.var("max_count", shape=(), dtype=iter_dtype), # iteration count + # iteration count + _expr.var("max_count", shape=(), dtype=iter_dtype), get_var(body.input[1].name, cond), # exit condition ] loop_vars += [get_var(body.input[i + 2].name, v) for i, v in enumerate(loop_deps)] @@ -4230,9 +4231,9 @@ def _impl_v10(cls, inputs, attr, params): dtype = infer_type(a).checked_type.dtype - ## Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 - ## and then requantize afer - ## https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/mlas/lib/qladd.cpp + # Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 + # and then requantize afer + # https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/mlas/lib/qladd.cpp a = _qnn.op.dequantize( inputs[0], a_scale, a_zero_point ) # , c_scale, c_zero_point, out_dtype = dtype) @@ -4296,7 +4297,8 @@ def try_resolve_to_const(x, dtype_override=None): b_zp_type = infer_type(b_zp).checked_type y_scale_type = infer_type(y_scale).checked_type - y_zp_type = infer_type(y_zp).checked_type # 'T3' in ONNX doc for this op + # 'T3' in ONNX doc for this op + y_zp_type = infer_type(y_zp).checked_type a_shape = infer_shape(a) b_shape = infer_shape(b) @@ -4471,9 +4473,9 @@ def _impl_v10(cls, inputs, attr, params): dtype = infer_type(a).checked_type.dtype - ## Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 - ## and then requantize afer - ## https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/mlas/lib/qlmul.cpp + # Onnxruntime doesn't actually do this op in integer, they dequantize to fp32 + # and then requantize afer + # https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/mlas/lib/qlmul.cpp a = _qnn.op.dequantize(inputs[0], a_scale, a_zero_point) b = _qnn.op.dequantize(inputs[3], b_scale, b_zero_point) out = _op.multiply(a, b) @@ -4515,10 +4517,10 @@ def _impl_v10(cls, inputs, attr, params): dtype = infer_type(x).checked_type.dtype - ## Apparently, onnxruntime doesn't do this op in integer, they dequantize to fp32 - ## and then requantize after: - ## https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/ - ## providers/dml/DmlExecutionProvider/src/GraphTransformer.cpp#L245 + # Apparently, onnxruntime doesn't do this op in integer, they dequantize to fp32 + # and then requantize after: + # https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/core/ + # providers/dml/DmlExecutionProvider/src/GraphTransformer.cpp#L245 x = _qnn.op.dequantize(x, x_scale, x_zero_point) out = _op.sigmoid(x) return _qnn.op.quantize(out, y_scale, y_zero_point, out_dtype=dtype) @@ -4663,12 +4665,16 @@ def _impl_v11(cls, inputs, attr, params): unique = _op.unique(data, is_sorted=(is_sorted == 1), return_counts=True) num_unique = unique[3] - trim_unique_lambda = lambda input: _op.strided_slice(input, _op.const([0]), num_unique) + def trim_unique_lambda(input): + return _op.strided_slice(input, _op.const([0]), num_unique) unique_vals = trim_unique_lambda(unique[0]) - indices = _op.cast(trim_unique_lambda(unique[1]), "int64") # ONNX always returns int64 - inverse_indices = _op.cast(unique[2], "int64") # ONNX always returns int64 - counts = _op.cast(trim_unique_lambda(unique[4]), "int64") # ONNX always returns int64 + # ONNX always returns int64 + indices = _op.cast(trim_unique_lambda(unique[1]), "int64") + # ONNX always returns int64 + inverse_indices = _op.cast(unique[2], "int64") + # ONNX always returns int64 + counts = _op.cast(trim_unique_lambda(unique[4]), "int64") # ONNX unique returns unique, indices, inverse_indices, (optional) counts return _expr.TupleWrapper(_expr.Tuple([unique_vals, indices, inverse_indices, counts]), 4) @@ -5087,6 +5093,37 @@ def _impl_v1(cls, inputs, attr, params): return _expr.TupleWrapper(_expr.Tuple(result), len(result)) +class Trilu(OnnxOpConverter): + """Operator converter for Trilu""" + + @classmethod + def _impl_v14(cls, inputs, attr, params): + upper = attr.get("upper", 1) + input_shape = shape_of(inputs[0]) + input_dims = infer_shape(input_shape)[0] + data_type = infer_type(inputs[0]).checked_type.dtype + k_tensor = relay.const(np.asarray(0), dtype=np.int64) + if len(inputs) == 2: + k_tensor = inputs[1] + + diag_input = relay.zeros(fold_constant(input_shape), dtype=data_type) + k1, k2 = None, None + if upper == 0: + k1 = relay.add(k_tensor, relay.const(1, dtype="int64")) + k1 = relay.expand_dims(k1, axis=0) + k2 = relay.take(input_shape, relay.const(input_dims - 1, dtype="int32")) + k2 = relay.expand_dims(k2, axis=0) + else: + k1 = relay.take(input_shape, relay.const(input_dims - 2, dtype="int32")) + k1 = relay.multiply(k1, relay.const(-1, dtype="int64")) + k1 = relay.subtract(k1, relay.const(1, dtype="int64")) + k1 = relay.expand_dims(k1, axis=0) + k2 = relay.subtract(k_tensor, relay.const(1, dtype="int64")) + k2 = relay.expand_dims(k2, axis=0) + + return relay.matrix_set_diag(inputs[0], diag_input, k=(k1, k2)) + + class Round(OnnxOpConverter): """Operator converter for round op.""" @@ -5114,6 +5151,8 @@ def _impl_v11(cls, inputs, attr, params): # use AttrCvt if attributes need to be converted # for 1 to N mapping(composed), use custom callable functions # for N to 1 mapping, currently not supported(?) + + def _get_convert_map(opset): return { # defs/experimental @@ -5287,6 +5326,7 @@ def _get_convert_map(opset): "CumSum": CumSum.get_converter(opset), "Unique": Unique.get_converter(opset), "Einsum": Einsum.get_converter(opset), + "Trilu": Trilu.get_converter(opset), # defs/control_flow "Loop": Loop.get_converter(opset), "If": If.get_converter(opset), @@ -5420,8 +5460,8 @@ def from_onnx(self, graph, opset, get_output_expr=False): # If requested, directly return the converted expressions. if get_output_expr: return outputs - ## Maintain the order of inputs and parameters from the ONNX graph, but only include - ## those parameters that are needed to execute the relay graph + # Maintain the order of inputs and parameters from the ONNX graph, but only include + # those parameters that are needed to execute the relay graph free_vars = analysis.free_vars(outputs) nodes = {v: k for k, v in self._nodes.items()} free_vars = [nodes[var] for var in free_vars] diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index 981074b6adb2..2278145312a8 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -535,7 +535,8 @@ def convert_qnn_fused_activation_function( raise ImportError("The tflite package must be installed") # Quantize a float value to an quantized integer value - quantize = lambda x: float(int(round(x / scale)) + zero_point) + def quantize(x): + return float(int(round(x / scale)) + zero_point) # Get min/max of the output dtype. This will be used to ensure that clip a_min/a_max are not # beyond the dtype range. @@ -1060,7 +1061,9 @@ def convert_relu_n1_to_1(self, op): # Quantize a float value to an quantized integer value scale_val = get_scalar_from_constant(input_tensor.qnn_params["scale"]) zero_point_val = get_scalar_from_constant(input_tensor.qnn_params["zero_point"]) - quantize = lambda x: float(int(round(x / scale_val)) + zero_point_val) + + def quantize(x): + return float(int(round(x / scale_val)) + zero_point_val) # Get min/max of the input dtype. This will be used to ensure that # clip a_min/a_max are not beyond the dtype range. @@ -3468,6 +3471,11 @@ def convert_matrix_set_diag(self, op): input_expr = self.get_tensor_expr(input_tensors[0]) diagonal_expr = self.get_tensor_expr(input_tensors[1]) + diag_shape = to_int_list(self.get_tensor_shape(input_tensors[1])) + input_shape = to_int_list(self.get_tensor_shape(input_tensors[0])) + if len(diag_shape) == len(input_shape) - 1: + diag_shape = np.insert(diag_shape, len(diag_shape) - 1, 1) + diagonal_expr = _op.reshape(diagonal_expr, diag_shape) out = _op.matrix_set_diag(input_expr, diagonal_expr) return out @@ -3488,13 +3496,13 @@ def convert_matrix_diag(self, op): scale and zero points to be equal" shape = to_int_list(self.get_tensor_shape(diagonal)) - shape = np.append(shape, shape[-1]) + diag_shape = np.insert(shape, len(shape) - 1, 1).astype(np.int32) dtype = self.get_tensor_type_str(diagonal.tensor.Type()) - + shape = np.append(shape, shape[-1]).astype(np.int32) input_expr = _op.zeros(tuple(shape), dtype) diagonal_expr = self.get_tensor_expr(diagonal) - out = _op.matrix_set_diag(input_expr, diagonal_expr) + out = _op.matrix_set_diag(input_expr, _op.reshape(diagonal_expr, diag_shape)) return out def convert_densify(self, op): diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index b5d44781e5e3..462bb33ee5a7 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -18,6 +18,7 @@ # pylint: disable=import-outside-toplevel """Transform operators.""" +import numpy as np from ...tir import expr as _expr from ..expr import Constant, Expr, Tuple, TupleWrapper, const from . import _make @@ -1247,7 +1248,7 @@ def sequence_mask(data, valid_length, mask_value=0, axis=0): def one_hot(indices, on_value, off_value, depth, axis, dtype): """ - Returns a one-hot tensor where the locations repsented by indices take value on_value, + Returns a one-hot tensor where the locations represented by indices take value on_value, other locations take value off_value. Final dimension is x depth x . @@ -1415,9 +1416,16 @@ def matrix_set_diag(data, diagonal, k=0, align="RIGHT_LEFT"): k_one = k k_two = k + if not isinstance(k_one, Expr): + k_one = const(np.asarray([k_one], dtype=np.int64)) + if not isinstance(k_two, Expr): + k_two = const(np.asarray([k_two], dtype=np.int64)) + super_diag_right_align = align[:5] == "RIGHT" sub_diag_right_align = align[-5:] == "RIGHT" + k_one = const(0) + k_two = const(0) return _make.matrix_set_diag( data, diagonal, k_one, k_two, super_diag_right_align, sub_diag_right_align ) diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py index d99d6772b0cd..6205cd495bf9 100644 --- a/python/tvm/topi/transform.py +++ b/python/tvm/topi/transform.py @@ -17,9 +17,12 @@ # pylint: disable=invalid-name,consider-using-enumerate,redefined-outer-name """Injective transformation operators""" from __future__ import absolute_import as _abs +import numpy as np +from tables import Expr import tvm from tvm import te from tvm import topi +from tvm.runtime.object_generic import const from tvm.te import hybrid from . import cpp from . import tag @@ -132,7 +135,7 @@ def flip(a, axis=0): The tensor to be expanded. axis : int, optional - The axis along which the tensors will be reveresed. + The axis along which the tensors will be reversed. Returns ------- @@ -183,7 +186,7 @@ def strided_slice(a, begin, end, strides=None, axes=None, slice_mode="end"): The indices to begin with in the slicing. end : list of int - Indicies indicating end of the slice. + Indices indicating end of the slice. strides : list of int, optional Specifies the stride values, it can be negative @@ -757,7 +760,7 @@ def where(condition, x, y): def one_hot(indices, on_value, off_value, depth, axis, dtype): """ - Returns a one-hot tensor where the locations repsented by indices take value on_value, + Returns a one-hot tensor where the locations represented by indices take value on_value, other locations take value off_value. Final dimension is x depth x . @@ -907,15 +910,33 @@ def matrix_set_diag(data, diagonal, k=0, align="RIGHT_LEFT"): [7, 5, 7, 7], [7, 7, 6, 7]]] """ + print("\n") + print("\n") + print("\n") if isinstance(k, (tuple, list)): + print("What is k 1 \n") k_one = k[0] if len(k) >= 2: k_two = k[1] else: k_two = k[0] else: + print("What is k 2 \n") k_one = k k_two = k + # k_one = te.placeholder(shape=(1,), name="k1", dtype="int64") + # k_two = te.placeholder(shape=(1,), name="k1", dtype="int64") + + # if not isinstance(k_one, Expr): + # k_one = const(np.asarray([k_one], dtype=np.int64)) + # if not isinstance(k_two, Expr): + # k_two = const(np.asarray([k_two], dtype=np.int64)) + + print(" one ", k_one) + print(" two ", k_two) + + # k_one = te.placeholder(shape=(1,), name="k1", dtype="int64") + # k_two = k_one super_diag_right_align = align[:5] == "RIGHT" sub_diag_right_align = align[-5:] == "RIGHT" diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 57bf9f36def9..e90984ea6913 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -3876,8 +3876,8 @@ TVM_REGISTER_NODE_TYPE(MatrixSetDiagAttrs); bool MatrixSetDiagRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { - // `types` contains: [input, diagonal, result] - ICHECK_EQ(types.size(), 3); + // `types` contains: [input, diagonal, k1, k2, result] + ICHECK_EQ(types.size(), 5); const auto* input = types[0].as(); ICHECK(input); @@ -3885,30 +3885,19 @@ bool MatrixSetDiagRel(const Array& types, int num_inputs, const Attrs& att const auto* diagonal = types[1].as(); ICHECK(diagonal); - const auto param = attrs.as(); - ICHECK_GE(param->k2, param->k1); + const auto* k1 = types[2].as(); + ICHECK(k1); - int d_ndims = diagonal->shape.size(); - int i_ndims = input->shape.size(); + const auto* k2 = types[3].as(); + ICHECK(k2); - reporter->Assert(input->shape[i_ndims - 2] > -param->k1); - reporter->Assert(input->shape[i_ndims - 1] > param->k2); + int d_ndims = diagonal->shape.size(); for (int i = 0; i < d_ndims - 2; i++) { reporter->AssertEQ(input->shape[i], diagonal->shape[i]); } - if (param->k1 != param->k2) { - reporter->AssertEQ(diagonal->shape[d_ndims - 2], param->k2 - param->k1 + 1); - } else if (d_ndims >= 2) { - reporter->AssertEQ(input->shape[d_ndims - 2], diagonal->shape[d_ndims - 2]); - } - auto max_diag_len = if_then_else(input->shape[i_ndims - 2] + (param->k2 > 0 ? param->k2 : 0) <= - input->shape[i_ndims - 1] + (param->k1 < 0 ? -param->k1 : 0), - input->shape[i_ndims - 2] + (param->k2 > 0 ? param->k2 : 0), - input->shape[i_ndims - 1] + (param->k1 < 0 ? -param->k1 : 0)); - reporter->AssertEQ(diagonal->shape[d_ndims - 1], max_diag_len); - reporter->Assign(types[2], TensorType(input->shape, input->dtype)); + reporter->Assign(types[4], TensorType(input->shape, input->dtype)); return true; } @@ -3916,20 +3905,21 @@ Array MatrixSetDiagCompute(const Attrs& attrs, const Array(); ICHECK(param != nullptr); - return Array{topi::matrix_set_diag(inputs[0], inputs[1], param->k1, param->k2, + std::cout << "**** \n" + << "d"; + printf("*******************\n"); + return Array{topi::matrix_set_diag(inputs[0], inputs[1], inputs[2], inputs[3], param->super_diag_right_align, param->sub_diag_right_align)}; } -Expr MakeMatrixSetDiag(Expr input, Expr diagonal, int k1, int k2, bool super_diag_right_align, +Expr MakeMatrixSetDiag(Expr input, Expr diagonal, Expr k1, Expr k2, bool super_diag_right_align, bool sub_diag_right_align) { auto attrs = make_object(); - attrs->k1 = k1; - attrs->k2 = k2; attrs->super_diag_right_align = super_diag_right_align; attrs->sub_diag_right_align = sub_diag_right_align; static const Op& op = Op::Get("matrix_set_diag"); - return Call(op, {input, diagonal}, Attrs(attrs), {}); + return Call(op, {input, diagonal, k1, k2}, Attrs(attrs), {}); } TVM_REGISTER_GLOBAL("relay.op._make.matrix_set_diag").set_body_typed(MakeMatrixSetDiag); @@ -3945,9 +3935,11 @@ RELAY_REGISTER_OP("matrix_set_diag") **sub_diag_right_align** Bool, true iff sub-diagonal is right aligned (left-padded). )code" TVM_ADD_FILELINE) .set_attrs_type() - .set_num_inputs(2) + .set_num_inputs(4) .add_argument("input", "Tensor", "Input Tensor.") .add_argument("diagonal", "Tensor", "Values to be filled in the diagonal.") + .add_argument("k1", "Tensor", "Lower limit (included) of the range of diagonals.") + .add_argument("k2", "Tensor", "Upper limit (included) of the range of diagonals.") .set_support_level(10) .add_type_rel("MatrixSetDiag", MatrixSetDiagRel) .set_attr("FTVMCompute", MatrixSetDiagCompute) diff --git a/src/topi/transform.cc b/src/topi/transform.cc index 56e799f52563..d6dc42237bd3 100644 --- a/src/topi/transform.cc +++ b/src/topi/transform.cc @@ -213,11 +213,10 @@ TVM_REGISTER_GLOBAL("topi.one_hot").set_body([](TVMArgs args, TVMRetValue* rv) { }); TVM_REGISTER_GLOBAL("topi.matrix_set_diag").set_body([](TVMArgs args, TVMRetValue* rv) { - int k1 = args[2]; - int k2 = args[3]; bool super_diag_right_align = args[4]; bool sub_diag_right_align = args[5]; - *rv = matrix_set_diag(args[0], args[1], k1, k2, super_diag_right_align, sub_diag_right_align); + *rv = matrix_set_diag(args[0], args[1], args[2], args[3], super_diag_right_align, + sub_diag_right_align); }); TVM_REGISTER_GLOBAL("topi.adv_index").set_body([](TVMArgs args, TVMRetValue* rv) { diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index c58e920ead1b..21dd9d63442b 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -2718,7 +2718,7 @@ def verify_conv( group=group, ) elif padding is None: - ## autopadding with unset default attributes + # autopadding with unset default attributes kwargs = {} if not all([s == 1 for s in strides]): kwargs["strides"] = strides @@ -3444,7 +3444,8 @@ def test_mod(target, dev): def verify_mod(x_shape, y_shape, fmod, out_shape, dtype="float32"): x_np = np.random.uniform(-100.0, 100.0, x_shape).astype(dtype) y_np = np.random.uniform(-100.0, 100.0, y_shape).astype(dtype) - y_np = np.where(y_np == 0, 1, y_np) # remove 0's to avoid division by zero error + # remove 0's to avoid division by zero error + y_np = np.where(y_np == 0, 1, y_np) mod_node = helper.make_node("Mod", inputs=["x", "y"], outputs=["z"], fmod=fmod) @@ -4353,11 +4354,13 @@ def verify_nonzero(indata, outdata, dtype): ) input_data = np.array([[1, 0], [1, 1]], dtype=np.int64) - result = np.array((np.nonzero(input_data))) # expected output [[0, 1, 1], [0, 0, 1]] + # expected output [[0, 1, 1], [0, 0, 1]] + result = np.array((np.nonzero(input_data))) verify_nonzero(input_data, result, dtype=np.int64) input_data = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]], dtype=np.int64) - result = np.array((np.nonzero(input_data))) # expected output [[0, 1, 2, 2], [0, 1, 0, 1]] + # expected output [[0, 1, 2, 2], [0, 1, 0, 1]] + result = np.array((np.nonzero(input_data))) verify_nonzero(input_data, result, dtype=np.int64) @@ -5224,24 +5227,6 @@ def verify_eyelike(indata): "test_training_dropout_mask", "test_training_dropout_zero_ratio", "test_training_dropout_zero_ratio_mask", - "test_tril", - "test_tril_pos", - "test_tril_square", - "test_tril_square_neg", - "test_tril_neg", - "test_tril_one_row_neg", - "test_tril_out_neg", - "test_tril_out_pos", - "test_tril_zero", - "test_triu", - "test_triu_one_row", - "test_triu_out_neg_out", - "test_triu_out_pos", - "test_triu_neg", - "test_triu_pos", - "test_triu_square", - "test_triu_square_neg", - "test_triu_zero", "test_unique_sorted_with_axis", "test_unique_sorted_with_axis_3d", "test_unique_sorted_with_negative_axis", @@ -5830,7 +5815,7 @@ def verify_qlinearconv( input_values.append(b_array) if padding is None: - ## autopadding with unset default attributes + # autopadding with unset default attributes kwargs = {} if not all([s == 1 for s in strides]): kwargs["strides"] = strides @@ -6357,7 +6342,7 @@ def verify_convinteger( input_values = [x_array, w_array] if padding is None: - ## autopadding with unset default attributes + # autopadding with unset default attributes kwargs = {} if not all([s == 1 for s in strides]): kwargs["strides"] = strides diff --git a/tests/python/relay/test_op_level10.py b/tests/python/relay/test_op_level10.py index a2d7f9938927..0b6dbd1bb8f5 100644 --- a/tests/python/relay/test_op_level10.py +++ b/tests/python/relay/test_op_level10.py @@ -676,7 +676,13 @@ def test_matrix_set_diag(executor_kind): def _verify(input_shape, diagonal_shape, dtype, k=0, align="RIGHT_LEFT"): input = relay.var("input", relay.TensorType(input_shape, dtype)) diagonal = relay.var("diagonal", relay.TensorType(diagonal_shape, dtype)) - out = relay.matrix_set_diag(input, diagonal, k, align) + + if len(diagonal_shape) == len(input_shape) - 1: + new_shape = list(diagonal_shape) + new_shape.insert(-1, 1) + out = relay.matrix_set_diag(input, relay.reshape(diagonal, new_shape), k, align) + else: + out = relay.matrix_set_diag(input, diagonal, k, align) in_type = run_infer_type(input) out_type = run_infer_type(out) diff --git a/tests/python/topi/python/test_topi_transform.py b/tests/python/topi/python/test_topi_transform.py index 180f267650cc..8279a4345a58 100644 --- a/tests/python/topi/python/test_topi_transform.py +++ b/tests/python/topi/python/test_topi_transform.py @@ -774,6 +774,88 @@ def check_device(target, dev): check_device(target, dev) +# def verify_matrix_set_diag(input_shape, diagonal_shape, dtype, k=0, align="RIGHT_LEFT"): +# # input matrix that contains diagonals to be replaced +# input = te.placeholder(shape=input_shape, name="input", dtype=dtype) +# # diagonal values to be placed as new diagonal values of input matrix +# diagonal = te.placeholder(shape=diagonal_shape, +# name="diagonal", dtype=dtype) +# # diagonals offsets +# # k1 and k2 define the lower and upper limits of diagonals to be set +# # where k*=0 means main diagonal, k*< 0 sub-diagonal, and k*> 0 super-diagonal +# # when k is not an tuple or list, k1 will be equal to k2, meaning that only one diagonal will be replaced. +# k1 = te.placeholder(shape=(1,), name="k1", dtype="int64") +# # k2 defines the upper limit diagonal to be set +# k2 = te.placeholder(shape=(1,), name="k2", dtype="int64") +# # matrix_set_diag_result = topi.transform.matrix_set_diag( +# # input, diagonal, (k1, k2), align) + +# matrix_set_diag_result = topi.transform.matrix_set_diag( +# input, diagonal, (k1, k2), align) +# `` +# # k can be an integer or a pair of integers representing the lower and upper limits of a matrix band; +# k_one, k_two = None, None +# if isinstance(k, (tuple, list)): +# print("define k *********") +# k_one = k[0] +# if len(k) >= 2: +# k_two = k[1] +# else: +# k_two = k[0] +# else: +# print("define k 12 *********") + +# k_one = k +# k_two = k + +# # Generate random data for input matrix +# input_npy = np.random.randint(-100, 100, size=input_shape).astype(dtype) +# # Generate random data for diagonal (single or multiple diagonals) +# diagonal_npy = np.random.randint(-100, 100, +# size=diagonal_shape).astype(dtype) +# # Run numpy test for matrix_set_diag with random data +# # output will be saved to compare with TOPI version of matrix_set_diag +# out_npy = tvm.topi.testing.matrix_set_diag( +# input_npy, diagonal_npy, k, align) + +# def check_device(target, dev): +# dev = tvm.device(target, 0) +# print("Running on target: %s" % target) +# with tvm.target.Target(target): +# s = tvm.topi.testing.get_injective_schedule( +# target)(matrix_set_diag_result) +# fn = tvm.build( +# s, [input, diagonal, k1, k2, +# matrix_set_diag_result], target, name="matrix_set_diag" +# ) + +# # Convert numpy input data to TVM ND array +# input_nd = tvm.nd.array(input_npy, dev) + +# # Convert numpy diagonal data to TVM ND array +# diagonal_nd = tvm.nd.array(diagonal_npy, dev) + +# # Convert k1 and k2 to numpy array and then to TVM ND array +# k1_nd = tvm.nd.array(np.asarray([k_one]), dev) +# k2_nd = tvm.nd.array(np.asarray([k_two]), dev) + +# # Convert k1 and k2 to numpy array and then to TVM ND array +# out_nd = tvm.nd.array(np.empty(out_npy.shape).astype( +# matrix_set_diag_result.dtype), dev) + +# # Run TOPI test for matrix_set_diag with random data +# fn(input_nd, diagonal_nd, k1_nd, k2_nd, out_nd) + +# # Convert TOPI output to numpy +# out_topi = out_nd.numpy() + +# # Check if Numpy version matches TOPI one +# tvm.testing.assert_allclose(out_topi, out_npy) + +# for target, dev in tvm.testing.enabled_targets(): +# check_device(target, dev) + + def verify_adv_index(data_shape, index_shapes, indice_dtype="int64"): dtype = "float32" data = te.placeholder(shape=data_shape, name="data", dtype=dtype) @@ -1226,7 +1308,8 @@ def test_sparse_to_dense(): verify_sparse_to_dense( [0, 1, 4], [3.1, 3.1, 3.1], 3.5, [5], [3.1, 3.1, 3.5, 3.5, 3.1] ) # floats - verify_sparse_to_dense(1, 3, None, [5], [0, 3, 0, 0, 0]) # default value not specified + # default value not specified + verify_sparse_to_dense(1, 3, None, [5], [0, 3, 0, 0, 0]) # negative test cases # sparse indices should be ints @@ -1241,8 +1324,8 @@ def test_sparse_to_dense(): def test_matrix_set_diag(): for dtype in ["float32", "int32"]: verify_matrix_set_diag((2, 2), (2,), dtype) - verify_matrix_set_diag((4, 3, 3), (4, 3), dtype) - verify_matrix_set_diag((2, 3, 4), (2, 3), dtype, 1) + # verify_matrix_set_diag((4, 3, 3), (4, 3), dtype) + # verify_matrix_set_diag((2, 3, 4), (2, 3), dtype, 1) verify_matrix_set_diag((2, 3, 4), (2, 4, 3), dtype, (-1, 2), "LEFT_RIGHT") verify_matrix_set_diag((2, 3, 4), (2, 4, 3), dtype, (-1, 2), "LEFT_LEFT") verify_matrix_set_diag((2, 3, 4), (2, 4, 3), dtype, (-1, 2), "RIGHT_RIGHT")