diff --git a/imperative/python/test/unit/jit/test_tracing.py b/imperative/python/test/unit/jit/test_tracing.py index 480587a925843212da548b2bf7252ec2af5299bd..6bc9e339ad4adb75e56a51e28cafbd492e8ec23d 100644 --- a/imperative/python/test/unit/jit/test_tracing.py +++ b/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], diff --git a/imperative/src/impl/physical_tensor.cpp b/imperative/src/impl/physical_tensor.cpp index 778a0fb11379ada6f550eee7532ba809472cdab9..5903b3dcb2bd00e0979055b4b791a97cd3466611 100644 --- a/imperative/src/impl/physical_tensor.cpp +++ b/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) diff --git a/src/core/include/megbrain/tensor.h b/src/core/include/megbrain/tensor.h index fe73be174f7ffda141491a9e80ad21e752f8fd0b..e22133d11e67bf7aa4fc1b428a7fae8983306524 100644 --- a/src/core/include/megbrain/tensor.h +++ b/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();