warp dim util

oneflow/api/python/framework/tensor_functions.cpp

@@ -22,6 +22,7 @@ limitations under the License.
 #include "oneflow/core/common/shape_vec.h"
 #include "oneflow/core/functional/functional.h"
 #include "oneflow/core/common/shape.h"
+#include "oneflow/core/common/wrap_dim_utils.h"
 
 namespace oneflow {
 namespace one {
@@ -329,10 +330,7 @@ static PyObject* PyTensorObject_size(PyObject* self, PyObject* args, PyObject* k
   if (idx_obj == NULL || idx_obj == Py_None) return TensorSize_NewFromShape(*shape);
   int64_t idx = PyLong_AsLongLong(idx_obj);
   int64_t ndim = shape->NumAxes();
-
-  CHECK_OR_THROW(idx >= -ndim && idx < ndim)
-      << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
-      << ", " << ndim - 1 << "], but got " << idx << ")";
+  idx = CHECK_JUST(maybe_wrap_dim(idx, ndim));
   idx = idx < 0 ? idx + ndim : idx;
   return PyLong_FromLongLong(shape->At(idx));
   END_HANDLE_ERRORS

oneflow/core/common/wrap_dim_utils.h

@@ -0,0 +1,40 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed 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,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/maybe.h"
+
+namespace oneflow {
+
+// align with pytorch: `c10/core/WrapDimMinimal.h`
+static inline Maybe<int64_t> maybe_wrap_dim(int64_t dim, int64_t dim_post_expr,
+                                            bool wrap_scalar = true) {
+  if (dim_post_expr <= 0) {
+    if (!wrap_scalar) {
+      return Error::RuntimeError()
+             << "dimension specified as " << dim << " but tensor has no dimensions";
+    }
+    dim_post_expr = 1;  // this will make range [-1, 0]
+  }
+
+  int64_t min = -dim_post_expr;
+  int64_t max = dim_post_expr - 1;
+  if (dim < min || dim > max) {
+    return Error::IndexError() << "Dimension out of range (expected to be in range of [" << min
+                               << ", " << max << "], but got " << dim << ")";
+  }
+  if (dim < 0) dim += dim_post_expr;
+  return dim;
+}
+}  // namespace oneflow

oneflow/core/framework/tensor_methods.cpp

@@ -24,6 +24,7 @@ limitations under the License.
 #include "oneflow/core/register/ofblob.h"
 #include "oneflow/core/framework/instructions_builder.h"
 #include "oneflow/core/ep/include/device_manager_registry.h"
+#include "oneflow/core/common/wrap_dim_utils.h"
 
 namespace oneflow {
 namespace one {
@@ -400,12 +401,7 @@ Maybe<Tensor> Transpose(const std::shared_ptr<Tensor>& input, const std::vector<
   CHECK_EQ_OR_RETURN(permute.size(), ndim)
       << "permute size should be equal to input tensor's ndim, but got " << permute.size();
   auto positive_perm = permute;
-  for (auto i = 0; i < positive_perm.size(); i++) {
-    if (positive_perm[i] < 0) { positive_perm[i] += ndim; }
-    CHECK_OR_RETURN(positive_perm[i] >= 0 && positive_perm[i] < ndim)
-        << "IndexError: Dimension out of range (expected to be in range of [" << -ndim << ","
-        << ndim << " ) but got " << positive_perm[i];
-  }
+  for (auto i = 0; i < positive_perm.size(); i++) { JUST(maybe_wrap_dim(positive_perm[i], ndim)); }
 
   DimVector target_dims(ndim);
   DimVector stride_vec(ndim);

oneflow/core/functional/function_library.h

@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_CORE_FUNCTIONAL_FUNCTION_LIBRARY_H_
 
 #include "oneflow/core/common/util.h"
+#include "oneflow/core/common/wrap_dim_utils.h"
 #include "oneflow/core/functional/packed_functor.h"
 #include "oneflow/core/common/stride.h"
 #include "oneflow/core/framework/tensor_methods.h"

oneflow/core/functional/impl/activation_functor.cpp

@@ -226,9 +226,7 @@ class GluFunctor {
     const auto ndim = input->ndim();
     CHECK_GT_OR_RETURN(ndim, 0) << Error::RuntimeError()
                                 << "glu does not support scalars because halving size must be even";
-    CHECK_OR_RETURN(dim >= -ndim && dim < ndim)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
-        << ", " << ndim - 1 << "], but got " << dim << ")";
+    dim = JUST(maybe_wrap_dim(dim, ndim));
     if (dim < 0) { dim += ndim; }
     int64_t nc = input->dim(dim);
     CHECK_EQ_OR_RETURN(nc % 2, 0) << Error::RuntimeError()
@@ -332,10 +330,7 @@ class SoftmaxFunctorBase {
     int64_t dim_ = dim ? JUST(dim) : get_dim();
     if (dim_ < 0) { dim_ += num_axes; }
 
-    CHECK_OR_RETURN(dim_ >= -num_axes && dim_ < num_axes)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of ["
-        << -num_axes << ", " << num_axes - 1 << "], but got " << dim_ << ")";
-
+    dim_ = JUST(maybe_wrap_dim(dim_, num_axes));
     if (dim_ != num_axes - 1) {
       std::vector<int> input_perm(input_shape->dim_vec().size(), 0);
       for (size_t i = 1; i < input_perm.size(); ++i) { input_perm[i] = i; }

oneflow/core/functional/impl/array_functor.cpp

@@ -68,9 +68,7 @@ class ArgMaxFunctor {
 
     int new_dim = JUST(dim);
     const int32_t ndims = input->shape()->NumAxes();
-    CHECK_OR_RETURN(new_dim >= -ndims && new_dim < ndims)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndims
-        << "," << ndims << " ] but got " << new_dim << ")";
+    new_dim = JUST(maybe_wrap_dim(new_dim, ndims));
     if (new_dim < 0) { new_dim += ndims; }
     const auto do_cast = [&](const std::shared_ptr<one::Tensor>& x) -> Maybe<Tensor> {
       return Cast(x, JUST(dtype), /*pin_memory=*/false);
@@ -469,10 +467,7 @@ class ConcatFunctor {
     int64_t ndim = inputs[0]->ndim();
     int64_t max_dim_size = 0;
     CHECK_GE_OR_RETURN(ninput, 1) << Error::RuntimeError() << "inputs size must greater than 0";
-    CHECK_OR_RETURN((-(ndim) <= dim) && (dim <= (ndim - 1)))
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
-        << ", " << ndim - 1 << "], but got " << dim << ")";
-    if (dim < 0) { axis += ndim; }
+    axis = JUST(maybe_wrap_dim(axis, ndim));
 
     const std::shared_ptr<const Shape>& shape = inputs[0]->shape();
     for (const auto& input : inputs) {
@@ -526,10 +521,7 @@ class StackFunctor {
     const int64_t ninput = inputs.size();
     int64_t ndims = inputs[0]->ndim();
     int64_t stack_dim = dim;
-    if (dim < 0) { stack_dim = stack_dim + ndims + 1; }
-    CHECK_OR_RETURN(stack_dim >= 0 && stack_dim <= ndims)
-        << Error::IndexError() << "Dimension out of range (expected in range of [" << -ndims - 1
-        << ", " << ndims << "], but got " << stack_dim << ")";
+    stack_dim = JUST(maybe_wrap_dim(stack_dim, ndims + 1));
     if (ninput == 1) { return ExpandDims(inputs[0], dim); }
     const std::shared_ptr<const Shape>& first_in_shape = inputs[0]->shape();
     for (const auto& input : inputs) {
@@ -666,9 +658,7 @@ class ExpandDimsFunctor {
   Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& input, const int32_t& dim) const {
     int32_t expand_dim = dim;
     const int32_t ndim = input->shape()->NumAxes();
-    CHECK_OR_RETURN(-(ndim + 1) <= dim && dim <= ndim)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of ["
-        << -(ndim + 1) << ", " << ndim << "], but got " << dim << ")";
+    JUST(maybe_wrap_dim(dim, ndim + 1));
     if (dim < 0) { expand_dim = dim + ndim + 1; }
     MutableAttrMap attrs;
     JUST(attrs.SetAttr<int32_t>("axis", expand_dim));
@@ -695,10 +685,7 @@ class SqueezeFunctor {
     if (dim.has_value()) {
       std::vector<int32_t> dims = *JUST(dim);
       for (int32_t dim_i : dims) {
-        CHECK_OR_RETURN((dim_i >= -ndim) && (dim_i <= ndim - 1))
-            << Error::IndexError() << "Dimension out of range (expected to be in range of  ["
-            << -ndim << "," << ndim - 1 << "], but got " << dim_i << ")";
-        if (dim_i < 0) { dim_i += ndim; }
+        dim_i = JUST(maybe_wrap_dim(dim_i, ndim));
         if (x->shape()->At(dim_i) == 1) { squeeze_dims.emplace_back(dim_i); }
       }
     } else {
@@ -776,16 +763,8 @@ class DimGatherFunctor {
         << Error::RuntimeError() << "gather(): Expected dtype int32 or int64 for index";
     CHECK_EQ_OR_RETURN(sparse_grad, false)
         << Error::RuntimeError() << "Only support bool = False for now!";
-    if (index->ndim() > 0) {
-      CHECK_LT_OR_RETURN(dim, index->ndim())
-          << Error::RuntimeError() << "Dimension out of range (expected to be in range of ["
-          << -index->ndim() << ", " << index->ndim() - 1 << "], but got " << dim << ")";
-    } else {
-      // For 0-dim Tensor
-      CHECK_LE_OR_RETURN(dim, index->ndim())
-          << Error::RuntimeError()
-          << "Dimension out of range (expected to be in range of [-1, 0], but got " << dim << ")";
-    }
+
+    JUST(maybe_wrap_dim(dim, index->ndim()));
     if (input->ndim() > 0 && index->ndim() > 0) {
       CHECK_EQ_OR_RETURN(input->ndim(), index->ndim())
           << Error::RuntimeError()
@@ -1300,11 +1279,8 @@ class NarrowFunctor {
     const int64_t ndim = input->shape()->NumAxes();
     CHECK_GT_OR_RETURN(ndim, 0) << Error::RuntimeError()
                                 << "narrow() cannot be applied to a 0-dim tensor.";
-    CHECK_OR_RETURN((-ndim <= dim) && (dim <= ndim - 1))
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
-        << ", " << ndim - 1 << "], but got " << dim << ")";
-    if (narrow_dim < 0) { narrow_dim += ndim; }
-    const int64_t dim_length = input->shape()->At(narrow_dim);
+    narrow_dim = JUST(maybe_wrap_dim(narrow_dim, ndim));
+    int64_t dim_length = input->shape()->At(narrow_dim);
     CHECK_OR_RETURN((-dim_length <= start) && (start <= dim_length))
         << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
         << ", " << ndim << "], but got " << start << ")";
@@ -1940,16 +1916,10 @@ class DiagonalFunctor {
   Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const int32_t& offset,
                            const int32_t& dim1, const int32_t& dim2) const {
     int64_t ndims = x->shape()->NumAxes();
-
-    CHECK_OR_RETURN(dim1 >= -ndims && dim1 < ndims)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndims
-        << ", " << ndims - 1 << "], but got " << dim1 << ")";
-    CHECK_OR_RETURN(dim2 >= -ndims && dim2 < ndims)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndims
-        << ", " << ndims - 1 << "], but got " << dim2 << ")";
-
-    const int32_t p_dim1 = dim1 >= 0 ? dim1 : dim1 + ndims;
-    const int32_t p_dim2 = dim2 >= 0 ? dim2 : dim2 + ndims;
+    int32_t p_dim1 = dim1;
+    int32_t p_dim2 = dim2;
+    p_dim1 = JUST(maybe_wrap_dim(p_dim1, ndims));
+    p_dim2 = JUST(maybe_wrap_dim(p_dim2, ndims));
     CHECK_NE_OR_RETURN(p_dim1, p_dim2)
         << Error::RuntimeError() << "diagonal dimensions cannot be identical " << dim1 << ", "
         << dim2;
@@ -2362,10 +2332,7 @@ class SplitFunctor {
   Maybe<TensorTuple> operator()(const std::shared_ptr<one::Tensor>& x,
                                 const int64_t& split_size_or_sections, const int64_t& dim) const {
     int64_t axis = dim;
-    if (axis < 0) { axis += x->ndim(); }
-    CHECK_OR_RETURN(axis >= 0 && axis < x->ndim())
-        << Error::IndexError() << "Dimension out of range (expected to be in range of ["
-        << -x->ndim() << ", " << x->ndim() - 1 << "], but got " << axis << ")";
+    axis = JUST(maybe_wrap_dim(axis, x->ndim()));
     CHECK_GE_OR_RETURN(split_size_or_sections, 0)
         << Error::RuntimeError() << "split expects split_size be non-negative, but got split_size="
         << split_size_or_sections;
@@ -2389,10 +2356,7 @@ class UnbindFunctor {
   Maybe<TensorTuple> operator()(const std::shared_ptr<one::Tensor>& x, const int64_t& dim) const {
     int32_t axis = dim;
     const int32_t ndim = x->ndim();
-    if (axis < 0) { axis += ndim; }
-    CHECK_OR_RETURN((dim >= -ndim) && (dim < ndim))
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
-        << "," << ndim - 1 << "], but got " << dim << ")";
+    axis = JUST(maybe_wrap_dim(axis, ndim));
     int32_t dim_size = x->shape()->At(axis);
     std::shared_ptr<TensorTuple> chunk_res = JUST(functional::Chunk(x, dim_size, axis));
     TensorTuple unbinds(dim_size);
@@ -2415,10 +2379,7 @@ class ChunkFunctor {
                               << "chunk expects at least a 1-dimensional tensor.";
     CHECK_OR_RETURN(chunks > 0) << Error::RuntimeError()
                                 << "chunk expects `chunks` to be greater than 0, got: " << chunks;
-    CHECK_OR_RETURN(-ndim <= dim && dim <= (ndim - 1))
-        << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
-        << ", " << ndim - 1 << "], but got " << dim << ")";
-    if (dim < 0) { infferd_dim += ndim; }
+    infferd_dim = JUST(maybe_wrap_dim(infferd_dim, ndim));
 
     const auto dim_size = x->shape()->At(infferd_dim);
     int64_t split_size = (dim_size + chunks - 1) / chunks;
@@ -2470,10 +2431,7 @@ class SplitWithSizeFunctor {
                                 const std::vector<int64_t>& split_size_or_sections,
                                 const int64_t& dim) const {
     int64_t axis = dim;
-    if (axis < 0) { axis += x->ndim(); }
-    CHECK_OR_RETURN(axis >= 0 && axis < x->ndim())
-        << Error::IndexError() << "Dimension out of range (expected to be in range of ["
-        << -x->ndim() << ", " << x->ndim() - 1 << "], but got " << axis << ")";
+    axis = JUST(maybe_wrap_dim(axis, x->ndim()));
     int64_t dim_size = x->shape()->At(axis);
     int64_t num_splits = split_size_or_sections.size();
     TensorTuple splits(num_splits);
@@ -2657,11 +2615,7 @@ class IndexSelectFunctor {
     CHECK_EQ_OR_RETURN(index_dtype_flag, true)
         << Error::RuntimeError() << "index_select(): Expected dtype int32 or int64 for index";
     int64_t new_dim = dim;
-    if (dim < 0) { new_dim += input_num_axes; }
-    CHECK_LE_OR_RETURN(new_dim, input_num_axes)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of ["
-        << -input_num_axes << ", " << input_num_axes - 1 << "], but got " << new_dim << ")";
-
+    new_dim = JUST(maybe_wrap_dim(new_dim, input_num_axes));
     return JUST(functional::Gather(input, index, new_dim));
   }
 };
@@ -2982,10 +2936,7 @@ class RepeatInterLeaveIntFunctor {
       int32_t dim_ = JUST(dim);
       const auto& input_shape = input->shape();
       const int64_t& num_axes = input_shape->NumAxes();
-      if (dim_ < 0) { dim_ += num_axes; }
-      CHECK_OR_RETURN(dim_ >= -num_axes && dim_ < num_axes)
-          << Error::IndexError() << "Dimension out of range (expected to be in range of ["
-          << -num_axes << ", " << num_axes - 1 << "], but got " << dim_ << ")";
+      dim_ = JUST(maybe_wrap_dim(dim_, num_axes));
       std::shared_ptr<one::Tensor> repeats_expand = JUST(
           Expand(JUST(Constant(Shape{1}, Scalar(repeats), DType::Int32(), JUST(input->device()))),
                  Shape{input->shape()->At(dim_)}));
@@ -3030,10 +2981,7 @@ class RepeatInterLeaveTensorFunctor {
     int32_t dim_ = dim;
     const auto& input_shape = input->shape();
     const int64_t& num_axes = input_shape->NumAxes();
-    if (dim_ < 0) { dim_ += num_axes; }
-    CHECK_OR_RETURN(dim_ >= -num_axes && dim_ < num_axes)
-        << Error::IndexError() << "Dimension out of range (expected to be in range of ["
-        << -num_axes << ", " << num_axes - 1 << "], but got " << dim_ << ")";
+    dim_ = JUST(maybe_wrap_dim(dim_, num_axes));
     CHECK_OR_RETURN(repeats_shape->At(0) == input->shape()->At(dim_))
         << Error::RuntimeError() << "repeats must have the same size as input along dim";
     std::shared_ptr<one::Tensor> cumsum = JUST(Cumsum(repeats, 0, DType::Int32()));

oneflow/core/functional/impl/common.cpp

@@ -15,6 +15,7 @@ limitations under the License.
 */
 #include "oneflow/core/functional/impl/common.h"
 #include "oneflow/core/autograd/autograd_mode.h"
+#include "oneflow/core/common/wrap_dim_utils.h"
 
 namespace oneflow {
 namespace one {
@@ -39,14 +40,7 @@ Maybe<std::vector<int32_t>> CheckAxis(const std::vector<int32_t>& axis, const in
     std::vector<int32_t> reduce_axis(naxis);
     std::vector<int32_t> axis_num(ndim);
     for (int32_t i = 0; i < naxis; i++) {
-      CHECK_OR_RETURN(axis[i] >= -ndim && axis[i] < ndim)
-          << Error::IndexError() << "Dimension out of range (expected to be in range of [" << -ndim
-          << ", " << ndim - 1 << "], but got " << axis[i] << ")";
-      if (axis[i] < 0) {
-        reduce_axis[i] = axis[i] + ndim;
-      } else {
-        reduce_axis[i] = axis[i];
-      }
+      reduce_axis[i] = JUST(maybe_wrap_dim(axis[i], ndim));
       axis_num[reduce_axis[i]]++;
       CHECK_OR_RETURN(axis_num[reduce_axis[i]] < 2)
           << Error::RuntimeError() << "dim " << reduce_axis[i]