From 20364a83fb585cbb59556e447e843e46b2df3b9d Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Mon, 18 Feb 2019 15:24:22 -0800 Subject: [PATCH] Update docs --- docs/api/python/topi.rst | 2 ++ docs/langref/relay_op.rst | 2 ++ python/tvm/relay/op/transform.py | 33 ++++++++++++++++++++ src/relay/op/tensor/transform.cc | 1 - tests/python/relay/test_op_level3.py | 45 ++++++++++++++-------------- topi/python/topi/transform.py | 22 ++++++++++++++ 6 files changed, 81 insertions(+), 24 deletions(-) diff --git a/docs/api/python/topi.rst b/docs/api/python/topi.rst index 856bad198e88e..ec5d600dab2bb 100644 --- a/docs/api/python/topi.rst +++ b/docs/api/python/topi.rst @@ -67,6 +67,7 @@ List of operators topi.not_equal topi.greater_equal topi.less_equal + topi.arange topi.image.resize @@ -123,6 +124,7 @@ topi .. autofunction:: topi.power .. autofunction:: topi.greater .. autofunction:: topi.less +.. autofunction:: topi.arange topi.nn ~~~~~~~ diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst index e1f38c61eb1f1..d58ba2e666219 100644 --- a/docs/langref/relay_op.rst +++ b/docs/langref/relay_op.rst @@ -95,6 +95,7 @@ This level enables additional math and transform operators. tvm.relay.full_like tvm.relay.cast tvm.relay.split + tvm.relay.arange **Level 4: Broadcast and Reductions** @@ -216,6 +217,7 @@ Level 3 Definitions .. autofunction:: tvm.relay.full_like .. autofunction:: tvm.relay.cast .. autofunction:: tvm.relay.split +.. autofunction:: tvm.relay.arange Level 4 Definitions diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 3716011cb1b34..8e8ca1097d2be 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -250,6 +250,39 @@ def full_like(data, fill_value): def arange(stop, start=None, step=1, dtype="float32"): + """Return evenly spaced values within a given interval. + + Warning: Undefined behavior when dtype is incompatible with start/stop/step. + It could lead to different results compared to numpy, MXNet, pytorch, etc. + + Parameters + ---------- + stop : tvm.Expr + Stop of interval. The interval does not include this value. + + start : tvm.Expr, optional + Start of interval. The interval includes this value. The default start + value is 0. + + step : tvm.Expr, optional + Spacing between values. The default step size is 1. + + dtype : str, optional + The target data type. + + Returns + ------- + result : relay.Expr + The resulting tensor. + + Examples + -------- + .. code-block:: python + + relay.arange(5) = [0, 1, 2, 3, 4] + relay.arange(1, 5) = [1, 2, 3, 4] + relay.arange(1, 5, 1.5) = [1, 2.5, 4] + """ if start is None: start = 0 else: diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc index 751371bcc4e98..c8024d30a5e5d 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -891,7 +891,6 @@ bool ArangeRel(const Array& types, const ArangeAttrs* param = attrs.as(); IndexExpr num_elem = tvm::cast(tvm::Int(32), tvm::ceil( tvm::cast(tvm::Float(32), param->stop - param->start) / param->step)); - num_elem = tvm::ir::Simplify(num_elem); reporter->Assign(types[0], TensorTypeNode::make({num_elem}, param->dtype)); return true; } diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py index e074bead7d688..8329ee24ad8b2 100644 --- a/tests/python/relay/test_op_level3.py +++ b/tests/python/relay/test_op_level3.py @@ -460,7 +460,6 @@ def test_infer_type_prelu(): def test_arange(): def verify_arange(start, stop, step): - print(start, stop, step) dtype = "float32" if start is None and step is None: x = relay.arange(stop) @@ -494,26 +493,26 @@ def verify_arange(start, stop, step): if __name__ == "__main__": - # test_cast() - # test_zeros_ones() - # test_unary_identity() - # test_clip() - # test_transpose_infer_type() - # test_transpose() - # test_reshape_infer_type() - # test_reshape() - # test_reshape_like_infer_type() - # test_reshape_like() - # test_take_infer_type() - # test_take() - # test_full_infer_type() - # test_full() - # test_full_like_infer_type() - # test_full_like() - # test_infer_type_leaky_relu() - # test_infer_type_prelu() - # test_squeeze() - # test_squeeze_infer_type() - # test_squeeze_bad_axes_infer_type() - # test_split_infer_type() + test_cast() + test_zeros_ones() + test_unary_identity() + test_clip() + test_transpose_infer_type() + test_transpose() + test_reshape_infer_type() + test_reshape() + test_reshape_like_infer_type() + test_reshape_like() + test_take_infer_type() + test_take() + test_full_infer_type() + test_full() + test_full_like_infer_type() + test_full_like() + test_infer_type_leaky_relu() + test_infer_type_prelu() + test_squeeze() + test_squeeze_infer_type() + test_squeeze_bad_axes_infer_type() + test_split_infer_type() test_arange() diff --git a/topi/python/topi/transform.py b/topi/python/topi/transform.py index 9ee1af7c5eaa1..2358639d08cc8 100644 --- a/topi/python/topi/transform.py +++ b/topi/python/topi/transform.py @@ -292,6 +292,28 @@ def tensordot(a, b, axes): def arange(stop, start=None, step=1, dtype="float32"): + """Creates a tensor with evenly spaced values within a given interval. + + Parameters + ---------- + stop : tvm.Expr + Stop of interval. The interval does not include this value. + + start : tvm.Expr, optional + Start of interval. The interval includes this value. The default start + value is 0. + + step : tvm.Expr, optional + Spacing between values. The default step size is 1. + + dtype : str, optional + The target data type. + + Returns + ------- + result : tvm.Tensor + The resulting tensor. + """ if start is None: start = 0 else: