add check non negative function

Oneflow-Inc · May 29, 2024 · 656879e · 656879e
1 parent ad0986a
commit 656879e
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/oneflow/api/python/functional/tensor_api.cpp b/oneflow/api/python/functional/tensor_api.cpp
@@ -203,10 +203,7 @@ class TensorWithShapeGenericCtorFunctor {
   Maybe<Tensor> operator()(const Shape& shape, const Symbol<DType>& dtype,
                            const Optional<Symbol<Device>>& device) const {
     // NOTE(chengcheng): flow.Tensor or flow.tensor ONLY created by EagerTensor now.
-    for (const auto& s : shape) {
-      CHECK_OR_THROW(s >= 0) << "Trying to create tensor with negative dimension " << s << ": "
-                             << shape;
-    }
+    JUST(CheckSizeNonNegative(shape));
     LazyMode::Guard lazy_mode_disabled_guard(/*is_enabled*/ false);
     Symbol<Device> device_;
     if (device) {

diff --git a/oneflow/core/functional/impl/common.cpp b/oneflow/core/functional/impl/common.cpp
@@ -204,6 +204,15 @@ Maybe<void> CheckInplaceShapeCanExpandTo(const Shape& shape, const Shape& expand
   return Maybe<void>::Ok();
 }
 
+Maybe<void> CheckSizeNonNegative(const Shape& shape) {
+  for (const auto& s : shape) {
+    CHECK_OR_THROW(s >= 0)
+        << "Trying to create tensor with negative dimension " << s << ": "
+        << shape;
+  }
+  return Maybe<void>::Ok();
+}
+
 Optional<Stride> ComputeStride(const Shape& shape, const Stride& stride,
                                const Shape& target_shape) {
   /*************************************************

diff --git a/oneflow/core/functional/impl/common.h b/oneflow/core/functional/impl/common.h
@@ -39,6 +39,7 @@ Maybe<void> CheckInplaceValid(const std::shared_ptr<Tensor>& x);
 Maybe<void> CheckInplaceCastValid(const std::shared_ptr<Tensor>& x,
                                   const std::shared_ptr<Tensor>& x_cast);
 Maybe<void> CheckInplaceShapeCanExpandTo(const Shape& shape, const Shape& expand_shape);
+Maybe<void> CheckSizeNonNegative(const Shape& shape);
 Optional<Stride> ComputeStride(const Shape& shape, const Stride& stride, const Shape& target_shape);
 Maybe<Shape> InferShapeUnspecifiedDim(const int64_t& elem_count, const Shape& shape);
 

diff --git a/oneflow/core/functional/impl/random_functor.cpp b/oneflow/core/functional/impl/random_functor.cpp
@@ -202,11 +202,7 @@ class RandFunctor {
         OF_UNIMPLEMENTED() << "Only support floating dtype in rand().";
       }
     }
-
-    for (const auto& s : shape) {
-      CHECK_OR_THROW(s >= 0) << "Trying to create tensor with negative dimension " << s << ": "
-                             << shape;
-    }
+    JUST(CheckSizeNonNegative(shape));
 
     auto gen = generator.value_or(JUST(one::DefaultAutoGenerator()));
     gen = JUST(GetGeneratorForLazyOrGlobal(gen, LazyMode::is_enabled(), NullOpt, NullOpt));