Fix(FromNumpy): fix bug in stride (#7042)

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
Oneflow-Inc · Dec 16, 2021 · c146dbb · c146dbb
1 parent bbd0d1d
commit c146dbb
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/oneflow/api/python/functional/tensor_api.cpp b/oneflow/api/python/functional/tensor_api.cpp
@@ -231,8 +231,19 @@ class LocalTensorSharedNumpyDataFunctor {
     DataType data_type = JUST(numpy::GetOFDataTypeFromNpArray(array));
     Symbol<Device> device = JUST(Device::New("cpu"));
     const npy_intp* stride_ptr = PyArray_STRIDES(array);
-    const auto stride = std::make_shared<Stride>(DimVector(stride_ptr, stride_ptr + dim));
-    auto tensor_meta = std::make_shared<MirroredTensorMeta>(shape, data_type, device, stride, 0);
+    // stride
+    auto strides_vec = DimVector(stride_ptr, stride_ptr + dim);
+    auto element_size_in_bytes = PyArray_ITEMSIZE(array);
+    // NumPy strides use bytes. OneFlow strides use element counts.
+    for (auto& stride : strides_vec) {
+      if (stride % element_size_in_bytes != 0) {
+        return Error::RuntimeError() << "given numpy array strides not a multiple of the element "
+                                        "byte size. Copy the numpy array to reallocate the memory.";
+      }
+      stride /= element_size_in_bytes;
+    }
+    const auto strides = std::make_shared<Stride>(strides_vec);
+    auto tensor_meta = std::make_shared<MirroredTensorMeta>(shape, data_type, device, strides, 0);
 
     // Build TensorBuffer
     const auto& Free = [obj](char* dptr) {

diff --git a/python/oneflow/test/modules/test_from_numpy.py b/python/oneflow/test/modules/test_from_numpy.py
@@ -28,6 +28,9 @@ def test_same_data(test_case):
         np_arr = np.random.randn(3, 4, 5)
         tensor = flow.from_numpy(np_arr)
         test_case.assertTrue(np.array_equal(np_arr, tensor.numpy()))
+        test_case.assertEqual(tensor.size(), (3, 4, 5))
+        test_case.assertEqual(tensor.stride(), (20, 5, 1))
+        test_case.assertEqual(tensor.storage_offset(), 0)
 
         np_arr[1:2, 2:3, 3:4] = random.random()
         test_case.assertTrue(np.array_equal(np_arr, tensor.numpy()))