[PyTorch] Fix shape and size() for columnwise-only quantized tensors

tests/pytorch/test_quantized_tensor.py

@@ -757,6 +757,58 @@ def test_shape_with_none_data(
             f"after setting data to None on {type(x_test).__name__}"
         )
 
+    @pytest.mark.parametrize("quantization", _quantization_list)
+    @pytest.mark.parametrize(
+        "rowwise, columnwise",
+        [(True, True), (True, False), (False, True)],
+        ids=["rowwise_columnwise", "rowwise_only", "columnwise_only"],
+    )
+    @pytest.mark.parametrize("shape", [(128, 256), (4, 128, 256)], ids=["2d", "3d"])
+    def test_shape_matches_size(
+        self,
+        *,
+        quantization: str,
+        rowwise: bool,
+        columnwise: bool,
+        shape: Iterable[int],
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+    ) -> None:
+        """shape, size() and size(dim) stay consistent for every usage combination.
+
+        Both shape and size() are derived from whichever data buffer is present,
+        and classes that store columnwise data transposed have to undo that. A
+        columnwise-only tensor is where they can drift apart -- from each other,
+        and from the shape the tensor was allocated with.
+        """
+        quantizer = make_quantizer(quantization, device=device)
+        # Row-scaled NVFP4 accepts set_usage(rowwise=False) but rejects the
+        # allocation itself, so it has to be filtered out up front.
+        if getattr(quantizer, "row_scaled_nvfp4", False) and not rowwise:
+            pytest.skip(f"{quantization} requires rowwise usage")
+        quantizer.set_usage(rowwise=rowwise, columnwise=columnwise)
+        if (quantizer.rowwise_usage, quantizer.columnwise_usage) != (rowwise, columnwise):
+            pytest.skip(f"{quantization} does not support this usage combination")
+
+        x = quantizer.make_empty(shape, dtype=dtype, device=device)
+        name = type(x).__name__
+
+        # shape and size() must describe the same tensor, whichever buffer they
+        # end up reading.
+        assert tuple(x.shape) == tuple(x.size()), f"{name}: {tuple(x.shape)} vs {tuple(x.size())}"
+
+        # size(dim) must agree with the full shape, including negative indices.
+        # It cannot be served by forwarding dim to a transposed buffer.
+        for dim in range(len(x.shape)):
+            assert x.size(dim) == x.shape[dim], f"{name}.size({dim}) is {x.size(dim)}"
+            neg = dim - len(x.shape)
+            assert x.size(neg) == x.shape[neg], f"{name}.size({neg}) is {x.size(neg)}"
+
+        # NVFP4 deliberately reports columnwise-only tensors flattened to 2D and
+        # warns about it, so only the ranks it preserves are checked here.
+        if not (isinstance(quantizer, NVFP4Quantizer) and not rowwise and len(shape) > 2):
+            assert tuple(x.shape) == tuple(shape), f"{name}.shape is {tuple(x.shape)}"
+
     @pytest.mark.parametrize(
         "quantization",
         _quantization_list + (["nvfp4_2d"] if nvfp4_available else []),

transformer_engine/pytorch/tensor/float8_blockwise_tensor.py

@@ -551,7 +551,11 @@ def shape(self):
         if self._rowwise_data is not None:
             return self._rowwise_data.shape
         if self._columnwise_data is not None:
-            return self._columnwise_data.shape
+            # Columnwise data is stored transposed, matching size() in the storage.
+            dims = self._columnwise_data.shape
+            if len(dims) == 2:
+                return torch.Size((dims[1], dims[0]))
+            return torch.Size(tuple(dims[1:]) + (dims[0],))
         return torch.Tensor.size(self)
 
     @property

transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py

@@ -317,12 +317,15 @@ def size(self, *args, **kwargs):
         # pylint: disable=missing-function-docstring
         if self._rowwise_data is not None:
             return self._rowwise_data.size(*args, **kwargs)
-        dims = list(self._columnwise_data.size(*args, **kwargs))
-        reordered = []
-        for i in range(1, len(dims)):
-            reordered.append(dims[i])
-        reordered.append(dims[0])
-        return torch.Size(reordered)
+        # Columnwise data is stored transposed, so a dim argument cannot be
+        # forwarded to it: rebuild the logical shape first, then index into it.
+        dims = self._columnwise_data.shape
+        if len(dims) == 2:
+            shape = torch.Size((dims[1], dims[0]))
+        else:
+            shape = torch.Size(tuple(dims[1:]) + (dims[0],))
+        dim = args[0] if args else kwargs.get("dim")
+        return shape if dim is None else shape[dim]
 
     def view(self, shape):
         """Reshape the leading (token) dims without dequantizing.

transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py

@@ -5,7 +5,6 @@
 """Mixin class holding data specific for Float8Tensor"""
 
 from __future__ import annotations
-import math
 from typing import Any, Dict, Optional, Tuple, Union
 import torch
 
@@ -179,8 +178,16 @@ def size(self, *args, **kwargs):
         # pylint: disable=missing-function-docstring
         if self._data is not None:
             return self._data.size(*args, **kwargs)
-        size = self._transpose.size(*args, **kwargs)
-        return torch.Size([size[-1], math.prod(size[:-1])])
+        # The transpose is stored as [last, *leading], so a dim argument cannot
+        # be forwarded to it: rebuild the logical shape first, then index into
+        # it. This matches the shape property on Float8Tensor.
+        dims = self._transpose.shape
+        if len(dims) == 2:
+            shape = torch.Size((dims[1], dims[0]))
+        else:
+            shape = torch.Size(tuple(dims[1:]) + (dims[0],))
+        dim = args[0] if args else kwargs.get("dim")
+        return shape if dim is None else shape[dim]
 
     @property
     def device(self):