This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 573d76f API Alias Redirect for torch aten tensor (#167)
573d76f is described below
commit 573d76f2e140f070497c1f61e384ed3720fa85cf
Author: Yaxing Cai <[email protected]>
AuthorDate: Fri Oct 17 16:03:18 2025 -0700
API Alias Redirect for torch aten tensor (#167)
This PR add the alias for `ndim -> dim`, `shape -> sizes` and
`IsContiguous -> is_contiguous` as torch aten tensor interface naming.
---
include/tvm/ffi/container/tensor.h | 60 +++++++++++++++++++++++++++++++++-----
tests/cpp/test_tensor.cc | 9 ++++++
2 files changed, 61 insertions(+), 8 deletions(-)
diff --git a/include/tvm/ffi/container/tensor.h
b/include/tvm/ffi/container/tensor.h
index 0b978db..da750e6 100644
--- a/include/tvm/ffi/container/tensor.h
+++ b/include/tvm/ffi/container/tensor.h
@@ -306,18 +306,26 @@ class Tensor : public ObjectRef {
}
/*!
- * \brief Get the size of the idx-th dimension.
+ * \brief Get the size of the idx-th dimension. If the idx is negative,
+ * it gets the size of last idx-th dimension.
* \param idx The index of the size.
* \return The size of the idx-th dimension.
*/
- int64_t size(size_t idx) const { return get()->shape[idx]; }
+ int64_t size(int64_t idx) const {
+ const TensorObj* ptr = get();
+ return ptr->shape[idx >= 0 ? idx : (ptr->ndim + idx)];
+ }
/*!
- * \brief Get the stride of the idx-th dimension.
+ * \brief Get the stride of the idx-th dimension. If the idx is negative,
+ * it gets the stride of last idx-th dimension.
* \param idx The index of the stride.
* \return The stride of the idx-th dimension.
*/
- int64_t stride(size_t idx) const { return get()->strides[idx]; }
+ int64_t stride(int64_t idx) const {
+ const TensorObj* ptr = get();
+ return ptr->strides[idx >= 0 ? idx : (ptr->ndim + idx)];
+ }
/*!
* \brief Get the number of elements in the Tensor.
@@ -473,6 +481,23 @@ class Tensor : public ObjectRef {
using ContainerType = TensorObj;
/// \endcond
+ // the following code are convenient APIs redirections created to provide
aten-style api
+ /*!
+ * \brief This functions redirects to ndim().
+ * \return The number of dimensions in the Tensor.
+ */
+ inline int32_t dim() { return ndim(); }
+ /*!
+ * \brief This functions redirects to shape().
+ * \return The shape of the Tensor.
+ */
+ inline ShapeView sizes() const { return shape(); }
+ /*!
+ * \brief This functions redirects to IsContiguous().
+ * \return True if the Tensor is contiguous, false otherwise.
+ */
+ inline bool is_contiguous() const { return IsContiguous(); }
+
protected:
/*!
* \brief Get const internal container pointer.
@@ -598,18 +623,20 @@ class TensorView {
}
/*!
- * \brief Get the size of the idx-th dimension.
+ * \brief Get the size of the idx-th dimension. If the idx is negative,
+ * it gets the size of last idx-th dimension.
* \param idx The index of the size.
* \return The size of the idx-th dimension.
*/
- int64_t size(size_t idx) const { return tensor_.shape[idx]; }
+ int64_t size(int64_t idx) const { return tensor_.shape[idx >= 0 ? idx :
tensor_.ndim + idx]; }
/*!
- * \brief Get the stride of the idx-th dimension.
+ * \brief Get the stride of the idx-th dimension. If the idx is negative,
+ * it gets the stride of last idx-th dimension.
* \param idx The index of the stride.
* \return The stride of the idx-th dimension.
*/
- int64_t stride(size_t idx) const { return tensor_.strides[idx]; }
+ int64_t stride(int64_t idx) const { return tensor_.strides[idx >= 0 ? idx :
tensor_.ndim + idx]; }
/*!
* \brief Get the byte offset of the Tensor.
@@ -623,6 +650,23 @@ class TensorView {
*/
bool IsContiguous() const { return tvm::ffi::IsContiguous(tensor_); }
+ // the following code are convenient APIs redirections created to provide
aten-style api
+ /*!
+ * \brief This functions redirects to ndim().
+ * \return The number of dimensions in the Tensor.
+ */
+ inline int32_t dim() { return ndim(); }
+ /*!
+ * \brief This functions redirects to shape().
+ * \return The shape of the Tensor.
+ */
+ inline ShapeView sizes() const { return shape(); }
+ /*!
+ * \brief This functions redirects to IsContiguous().
+ * \return True if the Tensor is contiguous, false otherwise.
+ */
+ inline bool is_contiguous() const { return IsContiguous(); }
+
private:
DLTensor tensor_;
template <typename, typename>
diff --git a/tests/cpp/test_tensor.cc b/tests/cpp/test_tensor.cc
index 41f7041..052a8ec 100644
--- a/tests/cpp/test_tensor.cc
+++ b/tests/cpp/test_tensor.cc
@@ -140,6 +140,15 @@ TEST(Tensor, EnvAlloc) {
EXPECT_EQ(tensor.size(0), 1);
EXPECT_EQ(tensor.size(1), 2);
EXPECT_EQ(tensor.size(2), 3);
+ EXPECT_EQ(tensor.size(-3), 1);
+ EXPECT_EQ(tensor.size(-2), 2);
+ EXPECT_EQ(tensor.size(-1), 3);
+ EXPECT_EQ(tensor.stride(0), 6);
+ EXPECT_EQ(tensor.stride(1), 3);
+ EXPECT_EQ(tensor.stride(2), 1);
+ EXPECT_EQ(tensor.stride(-3), 6);
+ EXPECT_EQ(tensor.stride(-2), 3);
+ EXPECT_EQ(tensor.stride(-1), 1);
EXPECT_EQ(tensor.dtype().code, kDLFloat);
EXPECT_EQ(tensor.dtype().bits, 32);
EXPECT_EQ(tensor.dtype().lanes, 1);
