support direct indexing for DenseTensor.dims() (PaddlePaddle#76331)

Author: HydrogenSulfate

paddle/phi/core/dense_tensor.h

@@ -85,6 +85,46 @@ class PADDLE_API DenseTensor : public TensorBase,
   /// \return The dims of the tensor.
   const DDim& dims() const noexcept override { return meta_.dims; }
 
+  /// \brief Returns the size of the tensor along the specified dimension.
+  ///        Supports negative indices, which count from the last dimension.
+  /// \param dim The dimension index to retrieve. Must be in the range [0, ndim)
+  /// or [-ndim, -1]. \return The size of the tensor along the given dimension.
+  /// \throws phi::errors::OutOfRange if the tensor is empty or the index is out
+  /// of range.
+  const int64_t dims(int dim) const {
+    int ndim = meta_.dims.size();
+
+    // Ensure the tensor has at least one dimension
+    PADDLE_ENFORCE_GE(ndim,
+                      1,
+                      phi::errors::OutOfRange(
+                          "dims expects at least a 1-dimensional tensor"));
+
+    // Check if the index is within the valid range [-ndim, ndim)
+    PADDLE_ENFORCE_GE(
+        dim,
+        -ndim,
+        phi::errors::OutOfRange(
+            "dims: dimension index (%d) must be in range [-%d, %d)",
+            dim,
+            ndim,
+            ndim));
+    PADDLE_ENFORCE_LT(
+        dim,
+        ndim,
+        phi::errors::OutOfRange(
+            "dims: dimension index (%d) must be in range [-%d, %d)",
+            dim,
+            ndim,
+            ndim));
+
+    // Handle negative indices
+    if (dim < 0) {
+      dim += ndim;
+    }
+    return meta_.dims[dim];
+  }
+
   /// \brief Returns the stride of the tensor.
   /// \return The stride of the tensor.
   const DDim& strides() const noexcept { return meta_.strides; }

test/cpp/phi/core/test_dense_tensor.cc

@@ -280,6 +280,56 @@ TEST(dense_tensor, shallow_copy) {
                         "tensor_1 to have the same meta"));
 }
 
+TEST(dense_tensor, dim_indexing) {
+  const DDim dims({4, 3, 2, 0});
+  const DataType dtype{DataType::INT8};
+  const DataLayout layout{DataLayout::NHWC};
+  const LegacyLoD lod{};
+  DenseTensorMeta meta(dtype, dims, layout, lod);
+
+  auto fancy_allocator = std::unique_ptr<Allocator>(new FancyAllocator);
+  auto* alloc = fancy_allocator.get();
+  DenseTensor tensor_0(alloc, meta);
+  int ndim = tensor_0.dims().size();
+  auto tensor_0_dims = tensor_0.dims();
+  for (int i = -ndim; i < ndim; ++i) {
+    PADDLE_ENFORCE_EQ(
+        tensor_0_dims[(i + ndim) % ndim],
+        tensor_0.dims(i),
+        common::errors::InvalidArgument(
+            "Dimension mismatch at index %d. Expected: %d, but got: %d",
+            i,
+            tensor_0_dims[i],
+            tensor_0.dims(i)));
+  }
+
+  // throw exception for index >= ndim
+  bool caught_exception = false;
+  try {
+    tensor_0.dims(ndim);
+  } catch (const common::enforce::EnforceNotMet& error) {
+    caught_exception = true;
+  }
+  PADDLE_ENFORCE_EQ(
+      caught_exception,
+      true,
+      common::errors::InvalidArgument(
+          "Expected an exception to be thrown for index >= ndim"));
+
+  // throw exception for index < -ndim
+  caught_exception = false;
+  try {
+    tensor_0.dims(-ndim - 1);
+  } catch (const common::enforce::EnforceNotMet& error) {
+    caught_exception = true;
+  }
+  PADDLE_ENFORCE_EQ(
+      caught_exception,
+      true,
+      common::errors::InvalidArgument(
+          "Expected an exception to be thrown for index < -ndim"));
+}
+
 TEST(dense_tensor, storage_properties) {
   const DataType dtype{DataType::FLOAT32};
   const DDim dims({1, 2});