diff --git a/docs/api/python/topi.rst b/docs/api/python/topi.rst index 12768cb7852f..9ac8bb1fd084 100644 --- a/docs/api/python/topi.rst +++ b/docs/api/python/topi.rst @@ -97,7 +97,7 @@ List of operators topi.repeat topi.tile topi.shape - topi.size + topi.ndarray_size topi.layout_transform topi.image.resize topi.argsort @@ -166,7 +166,7 @@ topi .. autofunction:: topi.repeat .. autofunction:: topi.tile .. autofunction:: topi.shape -.. autofunction:: topi.size +.. autofunction:: topi.ndarray_size .. autofunction:: topi.layout_transform .. autofunction:: topi.argsort .. autofunction:: topi.topk diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index 62aed04e01b1..e43fd5f7a2e7 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -288,10 +288,10 @@ struct SequenceMaskAttrs : public tvm::AttrsNode { }; // struct SequenceMaskAttrs. /*! \brief Attributes for ndarray_size operator */ -struct SizeAttrs : public tvm::AttrsNode { +struct NdarraySizeAttrs : public tvm::AttrsNode { DataType dtype; - TVM_DECLARE_ATTRS(SizeAttrs, "relay.attrs.SizeAttrs") { + TVM_DECLARE_ATTRS(NdarraySizeAttrs, "relay.attrs.NdarraySizeAttrs") { TVM_ATTR_FIELD(dtype) .describe("Target data type") .set_default(NullValue()); diff --git a/python/tvm/relay/op/contrib/contrib.py b/python/tvm/relay/op/contrib/contrib.py index c556f1f5bb0a..7114b7e712db 100644 --- a/python/tvm/relay/op/contrib/contrib.py +++ b/python/tvm/relay/op/contrib/contrib.py @@ -128,4 +128,4 @@ def ndarray_size(data, dtype="int32"): result : tvm.relay.Expr The number of elements of input tensor. """ - return _make.size(data, dtype) + return _make.ndarray_size(data, dtype) diff --git a/src/relay/op/tensor/unary.cc b/src/relay/op/tensor/unary.cc index b20cd8dfaaf4..60e53784649b 100644 --- a/src/relay/op/tensor/unary.cc +++ b/src/relay/op/tensor/unary.cc @@ -280,34 +280,34 @@ RELAY_REGISTER_OP("shape_of") .set_attr("FTVMCompute", ShapeOfCompute); -TVM_REGISTER_NODE_TYPE(NumElementsAttrs); +TVM_REGISTER_NODE_TYPE(NdarraySizeAttrs); -bool NumElementsRel(const Array& types, +bool NdarraySizeRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { CHECK_EQ(num_inputs, 1); auto tt = types[0].as(); CHECK(tt != nullptr); - const auto* param = attrs.as(); + const auto* param = attrs.as(); CHECK(param != nullptr); reporter->Assign(types[1], TensorTypeNode::make({1}, param->dtype)); return true; } -Array NumElementsCompute(const Attrs& attrs, +Array NdarraySizeCompute(const Attrs& attrs, const Array& inputs, const Type& out_type, const Target& target) { CHECK_EQ(inputs.size(), 1); - const auto* param = attrs.as(); + const auto* param = attrs.as(); CHECK(param != nullptr); - return Array{topi::size(inputs[0], param->dtype)}; + return Array{topi::ndarray_size(inputs[0], param->dtype)}; } -TVM_REGISTER_API("relay.op.contrib._make.size") +TVM_REGISTER_API("relay.op.contrib._make.ndarray_size") .set_body_typed([](Expr data, DataType dtype) { - auto attrs = make_node(); + auto attrs = make_node(); attrs->dtype = dtype; static const Op& op = Op::Get("contrib.ndarray_size"); return CallNode::make(op, {data}, Attrs(attrs), {}); @@ -318,15 +318,15 @@ RELAY_REGISTER_OP("contrib.ndarray_size") )code" TVM_ADD_FILELINE) .set_num_inputs(1) -.set_attrs_type_key("relay.attrs.NumElementsAttrs") +.set_attrs_type_key("relay.attrs.NdarraySizeAttrs") .add_argument("data", "Tensor", "The input tensor.") -.add_type_rel("NumElements", NumElementsRel) +.add_type_rel("NdarraySize", NdarraySizeRel) .set_attr("TOpIsStateful", false) .set_attr("TOpPattern", kInjective) .set_attr("FInferCorrectLayout", ElemwiseArbitraryLayout) .set_support_level(10) -.set_attr("FTVMCompute", NumElementsCompute); +.set_attr("FTVMCompute", NdarraySizeCompute); } // namespace relay } // namespace tvm diff --git a/tests/python/relay/test_op_level10.py b/tests/python/relay/test_op_level10.py index 52b6cc8023ce..f3520f3650a3 100644 --- a/tests/python/relay/test_op_level10.py +++ b/tests/python/relay/test_op_level10.py @@ -218,8 +218,8 @@ def test_shape_of(): def test_ndarray_size(): def verify_ndarray_size(shape): x = relay.var("x", shape=shape) - func = relay.Function([x], relay.op.contrib.num_elements(x)) - func = relay.ir_pass.infer_type(func) + func = relay.Function([x], relay.op.contrib.ndarray_size(x)) + func = run_infer_type(func) x_data = np.random.uniform(size=shape).astype("float32") ref_res = np.size(x_data) diff --git a/topi/include/topi/transform.h b/topi/include/topi/transform.h index b969fa8db2ed..c9a05098fc89 100644 --- a/topi/include/topi/transform.h +++ b/topi/include/topi/transform.h @@ -1231,13 +1231,13 @@ inline Tensor shape(const Tensor& src, * \param tag output tensor tag. * \return Tensor of input shape. */ -inline Tensor size(const Tensor& src, - const Type& dtype, - const std::string& name = "size", - const std::string& tag = kInjective) { +inline Tensor ndarray_size(const Tensor& src, + const Type& dtype, + const std::string& name = "ndarray_size", + const std::string& tag = kInjective) { int ndim = static_cast(src->shape.size()); - Array out_size = {1}; - return compute(out_size, [&](const Array& indices) { + Array out_ndarray_size = {1}; + return compute(out_ndarray_size, [&](const Array& indices) { Expr ret = 1; for (int i = 0; i < ndim; ++i) { ret *= src->shape[i]; diff --git a/topi/python/topi/transform.py b/topi/python/topi/transform.py index 8f3e034abe59..fc32403065d9 100644 --- a/topi/python/topi/transform.py +++ b/topi/python/topi/transform.py @@ -479,7 +479,7 @@ def sequence_mask(data, valid_length, mask_value=0, axis=0): return cpp.sequence_mask(data, valid_length, mask_value, axis) -def size(array, dtype="int32"): +def ndarray_size(array, dtype="int32"): """Get the number of elements of input array Parameters @@ -495,5 +495,4 @@ def size(array, dtype="int32"): result : tvm.Tensor The resulting tensor. """ - return cpp.size(array, dtype) - + return cpp.ndarray_size(array, dtype) diff --git a/topi/src/topi.cc b/topi/src/topi.cc index c8098cbf78b4..44134d7c2d67 100644 --- a/topi/src/topi.cc +++ b/topi/src/topi.cc @@ -311,9 +311,9 @@ TVM_REGISTER_GLOBAL("topi.shape") *rv = shape(args[0], args[1]); }); -TVM_REGISTER_GLOBAL("topi.size") +TVM_REGISTER_GLOBAL("topi.ndarray_size") .set_body([](TVMArgs args, TVMRetValue *rv) { - *rv = size(args[0], args[1]); + *rv = ndarray_size(args[0], args[1]); }); TVM_REGISTER_GLOBAL("topi.split") diff --git a/topi/tests/python/test_topi_transform.py b/topi/tests/python/test_topi_transform.py index 2d2eaf2b63a5..7f2c73e00390 100644 --- a/topi/tests/python/test_topi_transform.py +++ b/topi/tests/python/test_topi_transform.py @@ -649,11 +649,11 @@ def check_device(device): for backend in get_all_backend(): check_device(backend) -def test_size(): +def test_ndarray_size(): in_shape = (5, 11, 7) dtype = "int32" A = tvm.placeholder(shape=in_shape, dtype="float32", name="A") - B = topi.size(A, dtype) + B = topi.ndarray_size(A, dtype) input = np.random.uniform(size=in_shape).astype(A.dtype) output = np.asarray(np.size(input)).astype(dtype) @@ -668,7 +668,7 @@ def check_device(device): print("Running on target: %s" % device) with tvm.target.create(device): s = topi.generic.schedule_injective(B) - f = tvm.build(s, [A, B], device, name="size") + f = tvm.build(s, [A, B], device, name="ndarray_size") f(tvm_input, tvm_output) tvm.testing.assert_allclose(tvm_output.asnumpy(), output) @@ -695,4 +695,4 @@ def check_device(device): test_tile() test_shape() test_sequence_mask() - test_size() + test_ndarray_size()