fix(imperative/tensor): init m_offset when constructing a Tensor with DeviceTensorND

GitOrigin-RevId: b340e27c47

imperative/python/test/unit/jit/test_tracing.py

@@ -570,9 +570,9 @@ def test_random(shape_mode):
 def test_trace_advance_indexing(shape_mode):
     funcs = [
         lambda x, i: x[i],
-        # lambda x, i, j: x[i, j],  # FIXME
+        lambda x, i, j: x[i, j],
         lambda x, i, j: x[i, :, j, ...],
-        # lambda x, start, end: x[start:end],  # FIXME
+        lambda x, start, end: x[start:end],
         lambda x, start, end: x[:, 0, start:end, ..., 1],
         lambda x, vec: x[vec],
         lambda x, vec: x[vec, ..., 0, 1:3],

imperative/src/impl/physical_tensor.cpp

@@ -253,7 +253,7 @@ Tensor::Tensor(const DeviceTensorND &dv, const HostTensorND& hv) {
     }
     m_layout = dv.layout();
     m_blob = Blob::make(dv.storage());
-    m_offset = 0;
+    m_offset = dv.storage().offset();
 }
 
 Tensor::Tensor(const TensorLayout& layout, const CompNode& cn)

src/core/include/megbrain/tensor.h

@@ -176,6 +176,13 @@ class TensorStorage {
             return m_size;
         }
 
+        /*!
+         * \brief offset on allocated block in bytes
+         */
+        size_t offset() const {
+            return m_offset;
+        }
+
         //! get underlying comp node; error would be raised if it is invalid
         CompNode comp_node() const {
             check_comp_node_valid();