mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
97eb7a281d
The big semantic change (and the reason for this port) is that we no longer monkeypatch Tensor with torchdim's special methods. The new algorithm for handling dispatch is that we first land in `__torch_function__` and we see if a special FCD implementation needs to be dispatch to first, and if there is nothing we fallback to the standard level strategy. Because there is no longer C binding equivalent of classes, we've condensed _C.Dim and Dim together, and similar for Tensor. This resulted in some bugs as the Python API is sometimes different from the C API. I've attempted to disambiguate these but there may still be mistakes (many early bugs were due to this problem). Dim and DimEntry are especially painful as Dim must abide by Tensor equality semantics, but is pointer equality in C (DimEntry doesn't have this problem). Another difference between C/Python that is subtle is we no longer get implicit conversions from Dim to DimEntry, this also caused some bugs. Much of the mechanical porting work was done by claude code. I have a separate PR that deletes functorch._C, but it was useful having dim.cpp to point claude at it so I haven't done it in this PR. From a reviewing perspective, I need to re-review that I didn't forget to port anything, some noticeably missing "small" things are patched_dim_method. I am still in progress of carefully doing a side-by-side review of ports; "simplifications" from claude code were also a major source of bugs. There are two major feature gaps in the implementation: - DelayedTensor and dot handling are not implemented yet. This should be reasonably easy, just need to do it. However, for the purposes of sharded propagation it is actually better not to reconstruct matmuls. - Splitting dimensions with an index like `[x, y]` doesn't work. The problem is that `__getitem__` interprets this as advanced indexing and sends the list to torch.tensor to turn into a tensor, instead of being eligible for `__torch_function__`. I think I might need to hard code a special case for this or something? Signed-off-by: Edward Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/160236 Approved by: https://github.com/zdevito, https://github.com/albanD
69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Optional, TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from ._dim_entry import DimEntry
|
|
|
|
|
|
@dataclass
|
|
class TensorInfo:
|
|
tensor: Optional[torch.Tensor]
|
|
levels: list[DimEntry]
|
|
has_device: bool
|
|
batchedtensor: Optional[torch.Tensor]
|
|
|
|
def __post_init__(self) -> None:
|
|
from ._dim_entry import DimEntry
|
|
|
|
assert all(isinstance(l, DimEntry) for l in self.levels)
|
|
|
|
def ndim(self) -> int:
|
|
from ._dim_entry import ndim_of_levels
|
|
|
|
return ndim_of_levels(self.levels)
|
|
|
|
def __bool__(self) -> bool:
|
|
return self.tensor is not None
|
|
|
|
@staticmethod
|
|
def create(
|
|
h: Any, ensure_batched: bool = True, ensure_present: bool = True
|
|
) -> TensorInfo:
|
|
from . import Dim, DimEntry, Tensor
|
|
|
|
if Tensor.check_exact(h):
|
|
# functorch Tensor with first-class dimensions
|
|
return TensorInfo(
|
|
h._get_tensor(),
|
|
h._get_levels(),
|
|
h._get_has_device(),
|
|
h._get_batchtensor() if ensure_batched else None,
|
|
)
|
|
elif Dim.check_exact(h):
|
|
# For Dim objects, only get range/batchtensor if needed and dimension is bound
|
|
tensor = h._get_range() if h.is_bound else None
|
|
batchtensor = (
|
|
h._get_batchtensor() if ensure_batched and h.is_bound else None
|
|
)
|
|
return TensorInfo(
|
|
tensor,
|
|
[DimEntry(h)],
|
|
False,
|
|
batchtensor,
|
|
)
|
|
elif isinstance(h, torch.Tensor):
|
|
# Plain torch tensor - create positional levels
|
|
levels = []
|
|
for i in range(-h.dim(), 0):
|
|
levels.append(DimEntry(i))
|
|
return TensorInfo(h, levels, True, h)
|
|
else:
|
|
if ensure_present:
|
|
raise ValueError("expected a tensor object")
|
|
return TensorInfo(None, [], False, None)
|