pytorch/pytorch · error · ValueError

cannot preserve first-class dimensions with keepdim=True

Error message

cannot preserve first-class dimensions with keepdim=True

What it means

Thrown by `_wrap_dim` in functorch/dim `_wrap.py`. When a dim specification is a first-class `Dim` object and the caller requested `keepdim=True`, the reduction wrapper raises: a named first-class dimension cannot be 'kept' as a size-1 axis, because it would no longer index the original dim.

Source

Thrown at functorch/dim/_wrap.py:72

        def wrapped_func(*args: Any, **kwargs: Any) -> Any:
            return self.wrapper_implementation(self, *args, **kwargs)

        # Copy metadata using functools.update_wrapper for just __name__ and __doc__
        functools.update_wrapper(
            wrapped_func, self.orig, assigned=("__name__",), updated=()
        )
        wrapped_func.__doc__ = self.doc

        return wrapped_func


def _wrap_dim(dim: Any, ndim: int, keepdim: bool = False) -> DimEntry:
    """Convert single dimension specification to DimEntry object."""
    from . import Dim

    if isinstance(dim, Dim):
        if keepdim:
            raise ValueError("cannot preserve first-class dimensions with keepdim=True")
        return DimEntry(dim)
    elif isinstance(dim, int):
        i = dim
        while i >= 0:
            i -= ndim
        return DimEntry(i)
    else:
        return DimEntry()


def _wrap_dims(dim: Any, ndim: int, keepdim: bool = False) -> list[DimEntry]:
    """Convert dimension specification to list of DimEntry objects."""
    de = _wrap_dim(dim, ndim, keepdim)
    result = []
    if not de.is_none():
        result.append(de)
    else:
        for d in dim:

View on GitHub (pinned to dcd2ecae77)

Solutions

1. Drop `keepdim=True` when reducing over a Dim object; named dims handle broadcasting via the dim itself
2. If you need keepdim semantics, reduce over the positional int index (`t.sum(t.dims.index(d), keepdim=True)`) instead of the Dim
3. Refactor downstream code so it does not rely on the kept size-1 axis

Example fix

# before
s = t.sum(d, keepdim=True)  # d is a Dim -> ValueError

# after
s = t.sum(d)  # first-class dims do not support keepdim

from functorch.dim import Dim

def check_reduce_args(dims, keepdim):
    if keepdim and any(isinstance(d, Dim) for d in (dims if isinstance(dims, (list, tuple)) else [dims])):
        raise ValueError("first-class dims cannot be used with keepdim=True")

try:
    s = t.sum(d, keepdim=True)
except ValueError as e:
    if "cannot preserve first-class dimensions" in str(e):
        s = t.sum(d)  # retry without keepdim
    else:
        raise

When it happens

Trigger: Calling a wrapped reduction (e.g. `t.sum(d, keepdim=True)` in the dims API) where `d` is a `Dim` object rather than an int.

Common situations: Porting existing `keepdim=True` reduction code to first-class dims; passing dim objects from `dims('...')` while keeping the old keepdim habit.

Related errors

• Tensor with dimensions {level_strs} does not contain {d}
• Unable to get current frame
• Unable to get caller frame
• dims() must be assigned to a sequence of variable names or h
• expected {specified_ndims} sizes but found {len(sizes)}

AI-assisted analysis of pytorch/pytorch@dcd2ecae77 (2026-08-14). Data as JSON: /api/errors/2b3e3d08a68a3b51. Report an issue: GitHub.