[BUG] TensorDict.max(dim=..., reduce=True) and min(dim=..., reduce=True) raise AttributeError #2496

td.max(dim=0, reduce=True) and td.min(dim=0, reduce=True) raise AttributeError on any tensordict, so the documented reduce=True option works for these two reductions only without a dim. cummax and cummin with the same arguments return a torch.return_types result.

import torch
from tensordict import TensorDict

td = TensorDict(a=torch.arange(6.0).reshape(3, 2), batch_size=[3])
print(td.cummax(0, reduce=True).values)
print(type(td.max(dim=0, reduce=True, return_indices=False)).__name__)
td.max(dim=0, reduce=True)

tensor([[0., 1.],
        [2., 3.],
        [4., 5.]])
max
AttributeError: 'torch.return_types.max' object has no attribute 'keys'

Expected: td.max(dim=0, reduce=True) returns what torch.max gives on the leaves joined along dim 0, here torch.return_types.max(values=tensor([4., 5.]), indices=tensor([2, 2])), as cummax(0, reduce=True) already does. min should do the same. With return_indices=False it should return only the values tensor, as max(dim=0, return_indices=False) returns only the values tensordict.

Cause: with further_reduce=True, _cast_reduction returns the result of the torch reduction directly (tensordict/_base/reductions.py:2026-2036), which for max and min with a dim is a torch.return_types tuple. max (tensordict/_base/reductions.py:478-492) and min (tensordict/_base/reductions.py:304-318) then treat that result as a tensordict of values/indices pairs and call result.keys(...) on it. cummax and cummin return early when the result is a tensor or a torch.return_types (tensordict/_base/reductions.py:597 and :718); max and min lack that check. The reduce=True branch also ignores values_only, so return_indices=False does not drop the indices.

Possible fix: return early in max and min when the result is a tensor or a torch.return_types, as cummax and cummin do, and return .values from the reduce=True branch when values_only is set. max(dim=..., reduce=True) and min(dim=..., reduce=True) raise today. max(dim=..., reduce=True, return_indices=False), and the same call to min, cummax and cummin, return the torch.return_types tuple today and would return a tensor.

tensordict main (eba7879), torch 2.14.1, Python 3.12.12.