Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,17 @@
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import (
|
||||
DynamicOutputShapeException,
|
||||
FakeTensor,
|
||||
FakeTensorMode,
|
||||
UnsupportedFakeTensorException,
|
||||
)
|
||||
from torch._subclasses.fake_utils import CrossRefFakeMode
|
||||
|
||||
|
||||
__all__ = [
|
||||
"FakeTensor",
|
||||
"FakeTensorMode",
|
||||
"UnsupportedFakeTensorException",
|
||||
"DynamicOutputShapeException",
|
||||
"CrossRefFakeMode",
|
||||
]
|
||||
@@ -0,0 +1,273 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch import SymInt
|
||||
from torch.fx.experimental.sym_node import SymNode
|
||||
from torch.types import py_sym_types, PySymType
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import sympy
|
||||
|
||||
from torch.fx.experimental.symbolic_shapes import ShapeEnv
|
||||
|
||||
from .fake_tensor import _DispatchCacheKey, _MetadataIntLike
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _DeconstructedSymNode:
|
||||
"""
|
||||
Represents a SymNode without the associated ShapeEnv
|
||||
"""
|
||||
|
||||
# n.b. keep the same protocol as SymNode
|
||||
_expr: sympy.Expr
|
||||
pytype: type
|
||||
_hint: int | float | bool | None
|
||||
constant: int | float | bool | None
|
||||
fx_node: torch.fx.Node
|
||||
|
||||
@staticmethod
|
||||
def from_node(node: SymNode) -> _DeconstructedSymNode:
|
||||
return _DeconstructedSymNode(
|
||||
node._expr,
|
||||
node.pytype,
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
node._hint,
|
||||
node.constant,
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
node.fx_node,
|
||||
)
|
||||
|
||||
def extract(self, shape_env: ShapeEnv) -> SymNode:
|
||||
return SymNode(
|
||||
self._expr, shape_env, self.pytype, self._hint, self.constant, self.fx_node
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return str(self._expr)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"_DeconstructedSymNode{{{self._expr!r}, {self.pytype!r}, {self._hint!r}, {self.constant!r}, {self.fx_node!r}}}"
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def __hash__(self) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
# _value_eq to match SymNode
|
||||
def _value_eq(self, other: object) -> bool:
|
||||
if isinstance(other, (SymNode, _DeconstructedSymNode)):
|
||||
return (
|
||||
self._expr == other._expr
|
||||
and self.pytype == other.pytype
|
||||
and self._hint == other._hint
|
||||
and self.constant == other.constant
|
||||
and self.fx_node == other.fx_node
|
||||
)
|
||||
else:
|
||||
return False
|
||||
|
||||
# _value_hash to match SymNode
|
||||
def _value_hash(self) -> int:
|
||||
return hash((self._expr, self.pytype, self._hint, self.constant, self.fx_node))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _DeconstructedSymType:
|
||||
"""
|
||||
Represents a SymInt, SymFloat, SymBool without the associated ShapeEnv
|
||||
"""
|
||||
|
||||
ty: type[PySymType]
|
||||
node: _DeconstructedSymNode
|
||||
|
||||
@staticmethod
|
||||
def from_sym_type(value: PySymType) -> _DeconstructedSymType:
|
||||
return _DeconstructedSymType(type(value), value.node)
|
||||
|
||||
def extract(self, shape_env: ShapeEnv) -> PySymType:
|
||||
return self.ty(self.node.extract(shape_env))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.ty}({self.node})"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"_DeconstructedSymType({self.ty}, {self.node!r})"
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return NotImplemented
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return NotImplemented
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _InputBackref:
|
||||
value: int
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _PySymInputStub:
|
||||
"""
|
||||
Represents a SymInt in the cached key. Needed because SymInt doesn't
|
||||
support __eq__ or __hash__ directly.
|
||||
"""
|
||||
|
||||
# value can be:
|
||||
# PySymType: This is the 'normal' SymInt value, wrapped so we can use
|
||||
# hash/eq as value hash/eq (normally SymInt does object
|
||||
# hash/eq).
|
||||
# _DeconstructedSymType: This is used when storing the _PySymInputStub in
|
||||
# the cache to avoid cyclic ShapeEnv references.
|
||||
# _InputBackref: This is a back-reference to a previous _PySymInputStub in
|
||||
# the key.
|
||||
value: PySymType | _DeconstructedSymType | _InputBackref
|
||||
|
||||
def __init__(
|
||||
self, value: PySymType | _DeconstructedSymType | _InputBackref
|
||||
) -> None:
|
||||
# For inputs (values in the `key`) we need to keep the PySymType intact
|
||||
# - this way if we need to reuse it as an output we can properly copy
|
||||
# the original value.
|
||||
self.value = value
|
||||
|
||||
def strip_shape_env(self) -> None:
|
||||
if isinstance(self.value, py_sym_types):
|
||||
self.value = _DeconstructedSymType.from_sym_type(self.value)
|
||||
|
||||
def extract(self, shape_env: ShapeEnv) -> PySymType:
|
||||
if isinstance(self.value, _DeconstructedSymType):
|
||||
return self.value.extract(shape_env)
|
||||
else:
|
||||
# We should never see an _InputBackref here - anyone extracting a
|
||||
# value should be pulling from the original entry (the one this
|
||||
# backref points at).
|
||||
if isinstance(self.value, _InputBackref):
|
||||
raise AssertionError(
|
||||
"Cannot extract value from _InputBackref - use the original entry"
|
||||
)
|
||||
return self.value
|
||||
|
||||
def __str__(self) -> str:
|
||||
return str(self.value)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"_PySymInputStub({self.value!r})"
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, _PySymInputStub):
|
||||
return False
|
||||
elif isinstance(self.value, _InputBackref) or isinstance(
|
||||
other.value, _InputBackref
|
||||
):
|
||||
return self.value == other.value
|
||||
else:
|
||||
return self.value.node._value_eq(other.value.node)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
if isinstance(self.value, _InputBackref):
|
||||
return hash(self.value)
|
||||
else:
|
||||
return self.value.node._value_hash()
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _SymIntOutputStub:
|
||||
"""
|
||||
Represents a SymInt in the cached output.
|
||||
"""
|
||||
|
||||
# This is either an `int` which represents the index in the key to copy the
|
||||
# SymNode from or it's the deconstructed SymNode itself.
|
||||
value: int | _DeconstructedSymNode
|
||||
|
||||
def __init__(self, value: SymInt, key_path: int | None) -> None:
|
||||
if key_path is None:
|
||||
self.value = _DeconstructedSymNode.from_node(value.node)
|
||||
else:
|
||||
self.value = key_path
|
||||
|
||||
def extract(self, key: _DispatchCacheKey, shape_env: ShapeEnv) -> SymInt:
|
||||
if isinstance(self.value, _DeconstructedSymNode):
|
||||
return SymInt(self.value.extract(shape_env))
|
||||
else:
|
||||
src = key.key[self.value]
|
||||
if not isinstance(src, _PySymInputStub) or not isinstance(
|
||||
src.value, SymInt
|
||||
):
|
||||
raise AssertionError(
|
||||
f"Expected _PySymInputStub with SymInt value, got {type(src)} "
|
||||
f"with {type(src.value) if isinstance(src, _PySymInputStub) else 'N/A'}"
|
||||
)
|
||||
return src.value
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"_SymIntOutputStub({self.value!r})"
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
def __hash__(self) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _CacheKeyState:
|
||||
"""
|
||||
State used while building our cache key.
|
||||
"""
|
||||
|
||||
# We track the SymNodes so when we get the output we can see if it exactly
|
||||
# matches one of the inputs so we can uncache it properly.
|
||||
sym_node_lookup: dict[int, int] # id(SymNode) -> index
|
||||
|
||||
# This is a list of all seen input sympy.Symbols. We use it when building
|
||||
# the cache entry to see if the output value has any symbols that we didn't
|
||||
# see on input. See _has_unrepresented_symbols().
|
||||
known_symbols: set[sympy.Symbol]
|
||||
|
||||
# There are cases where we're asked to perform an op when we have no
|
||||
# ShapeEnv on the FakeTensorMode - but for SymNodes we MUST have a
|
||||
# ShapeEnv. So as we scan if we see a SymNode (with a ShapeEnv) we record it
|
||||
# here.
|
||||
shape_env: ShapeEnv | None
|
||||
|
||||
def __init__(self, shape_env: ShapeEnv | None = None) -> None:
|
||||
self.sym_node_lookup = {}
|
||||
self.known_symbols = set()
|
||||
self.shape_env = shape_env
|
||||
|
||||
def cache_on_shape_env(self) -> bool:
|
||||
"""
|
||||
Returns true if the CacheKey needs to be cached on the ShapeEnv
|
||||
rather than the global cache.
|
||||
|
||||
If our inputs contain a SymNode then we can't cache this operation on
|
||||
the global cache because the cached output will implicitly depend on
|
||||
guard values which might not be true on some other ShapeEnv. So unless
|
||||
we're also going to cache the guards we need to cache this operation on
|
||||
the ShapeEnv instead of globally.
|
||||
"""
|
||||
return bool(self.sym_node_lookup)
|
||||
|
||||
def convert_sym_int(self, result: list[object], arg: SymInt) -> None:
|
||||
node_id = id(arg.node)
|
||||
if node_id in self.sym_node_lookup:
|
||||
result.append(_InputBackref(self.sym_node_lookup[node_id]))
|
||||
else:
|
||||
self.sym_node_lookup[node_id] = len(result)
|
||||
self.known_symbols.update(arg.node.expr.free_symbols)
|
||||
if self.shape_env is None:
|
||||
self.shape_env = arg.node.shape_env
|
||||
result.append(_PySymInputStub(arg))
|
||||
|
||||
def convert_output(self, arg: _MetadataIntLike) -> _MetadataIntLike:
|
||||
if isinstance(arg, SymInt):
|
||||
return _SymIntOutputStub(arg, self.sym_node_lookup.get(id(arg.node), None))
|
||||
else:
|
||||
return arg
|
||||
@@ -0,0 +1,9 @@
|
||||
from ._core import ComplexTensor
|
||||
from ._ops import ComplexTensorMode, is_complex_tensor
|
||||
|
||||
|
||||
__all__ = ["ComplexTensor", "ComplexTensorMode", "is_complex_tensor"]
|
||||
|
||||
ComplexTensor.__module__ = __name__
|
||||
ComplexTensorMode.__module__ = __name__
|
||||
is_complex_tensor.__module__ = __name__
|
||||
@@ -0,0 +1,153 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing_extensions import Self
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.autograd import Function
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._ops import OpOverload
|
||||
from torch._prims_common import DeviceLikeType
|
||||
from torch.autograd.function import FunctionCtx
|
||||
|
||||
|
||||
class ComplexTensor(Tensor):
|
||||
"""A class that decomposes all ops on complex Tensors into their real and imaginary parts."""
|
||||
|
||||
_re: Tensor
|
||||
_im: Tensor
|
||||
|
||||
def __new__(cls, real: Tensor, imag: Tensor) -> Self:
|
||||
"""Initialize a ComplexTensor from its real and imaginary parts."""
|
||||
from ._ops.common import REAL_TO_COMPLEX
|
||||
|
||||
shape = real.shape
|
||||
device = real.device
|
||||
|
||||
# TODO (hameerabbasi): `torch.compile` sometimes fails here without making these
|
||||
# contiguous. Why?
|
||||
real = real.contiguous()
|
||||
imag = imag.contiguous()
|
||||
|
||||
# TODO (hameerabbasi):
|
||||
# What should we do with dtype?
|
||||
# We could convert to the complex type (float32 -> complex64), but we
|
||||
# can't use that model for say `bfloat16` which does not have a
|
||||
# corresponding complex dtype.
|
||||
# If we want to support this complex rep using any float type (see
|
||||
# https://github.com/pytorch/pytorch/issues/95100)
|
||||
# We either need to:
|
||||
# 1) add the complex types for say `complexbf32`, knowing they can't really be used anywhere
|
||||
# else.
|
||||
# 2) We use the real float dtype here, and it is up to the user to know
|
||||
# that dtype=float<size> here really means complex<2xSize> with dtype
|
||||
# matching that of re/im parts alone
|
||||
# I'm going with 1 for now, so that I can make gradcheck and some complex
|
||||
# ops work properly, but might want to discuss this in the RFP.
|
||||
dtype = REAL_TO_COMPLEX.get(real.dtype)
|
||||
if dtype is None:
|
||||
raise TypeError(
|
||||
"Unsupported dtype for constituent tensors. Supported dtypes are: "
|
||||
f"{set(REAL_TO_COMPLEX.keys())!r}."
|
||||
)
|
||||
storage_offset = real.storage_offset()
|
||||
strides = real.stride()
|
||||
layout = real.layout
|
||||
pin_memory = real.is_pinned()
|
||||
|
||||
if shape != imag.shape:
|
||||
raise AssertionError(f"Expected imag shape {shape}, got {imag.shape}")
|
||||
if device != imag.device:
|
||||
raise AssertionError(f"Expected imag device {device}, got {imag.device}")
|
||||
if real.dtype != imag.dtype:
|
||||
raise AssertionError(f"Expected imag dtype {real.dtype}, got {imag.dtype}")
|
||||
if pin_memory != imag.is_pinned():
|
||||
raise AssertionError(
|
||||
f"Expected imag pinning {pin_memory}, got {imag.is_pinned()}"
|
||||
)
|
||||
|
||||
res = Tensor._make_wrapper_subclass( # type: ignore[attr-defined]
|
||||
cls,
|
||||
shape,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
storage_offset=storage_offset,
|
||||
strides=strides,
|
||||
pin_memory=pin_memory,
|
||||
layout=layout,
|
||||
requires_grad=False,
|
||||
)
|
||||
res._re = real.clone().detach()
|
||||
res._im = imag.clone().detach()
|
||||
|
||||
return res
|
||||
|
||||
@property
|
||||
def re(self) -> Tensor:
|
||||
return self._re
|
||||
|
||||
@property
|
||||
def im(self) -> Tensor:
|
||||
return self._im
|
||||
|
||||
@classmethod
|
||||
def __torch_dispatch__( # type: ignore[bad-override]
|
||||
cls,
|
||||
func: OpOverload,
|
||||
types: tuple[type, ...],
|
||||
# pyrefly: ignore [implicit-any]
|
||||
args: tuple = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
from ._ops.common import lookup_complex
|
||||
|
||||
kwargs = {} if kwargs is None else kwargs
|
||||
|
||||
impl = lookup_complex(func, *args, **kwargs)
|
||||
if impl is None:
|
||||
return NotImplemented
|
||||
|
||||
return impl(*args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def from_interleaved(t: Tensor) -> ComplexTensor:
|
||||
t_real = torch.real(t)
|
||||
t_imag = torch.imag(t) if t.dtype.is_complex else torch.zeros_like(t_real)
|
||||
return Complex.apply(t_real, t_imag)
|
||||
|
||||
def as_interleaved(self) -> Tensor:
|
||||
return torch.complex(self.real, self.imag)
|
||||
|
||||
@staticmethod
|
||||
def __tensor_unflatten__(
|
||||
inner_tensors: dict[str, Tensor],
|
||||
meta: Any,
|
||||
outer_size: tuple[int, ...],
|
||||
outer_stride: tuple[int, ...],
|
||||
) -> ComplexTensor:
|
||||
if meta is not None:
|
||||
raise AssertionError(f"meta must be None, got {meta}")
|
||||
re, im = inner_tensors["re"], inner_tensors["im"]
|
||||
return ComplexTensor(re, im)
|
||||
|
||||
def __tensor_flatten__(self) -> tuple[list[str], Any]:
|
||||
return ["re", "im"], None
|
||||
|
||||
def __repr__(self, *, tensor_contents: object | None = None) -> str:
|
||||
return f"ComplexTensor(real={self.re!r}, imag={self.im!r})"
|
||||
|
||||
def is_pinned(self, device: DeviceLikeType | None = None) -> bool:
|
||||
return self.re.is_pinned(device)
|
||||
|
||||
|
||||
class Complex(Function):
|
||||
@staticmethod
|
||||
def forward(ctx: FunctionCtx, real: Tensor, imag: Tensor) -> ComplexTensor: # type: ignore[bad-override]
|
||||
return ComplexTensor(real, imag)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx: FunctionCtx, grad_output: ComplexTensor) -> tuple[Tensor, Tensor]: # type: ignore[bad-override]
|
||||
return grad_output.real, grad_output.imag
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
from . import aten, prims
|
||||
from .common import ComplexTensorMode, is_complex_tensor
|
||||
|
||||
|
||||
__all__ = ["ComplexTensorMode", "is_complex_tensor", "aten", "prims"]
|
||||
+971
@@ -0,0 +1,971 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from .._core import ComplexTensor
|
||||
from .common import (
|
||||
_get_func_name,
|
||||
COMPLEX_TO_REAL,
|
||||
complex_to_real_dtype,
|
||||
is_complex,
|
||||
OpType,
|
||||
promote_tensors,
|
||||
register_binary_nonlinear,
|
||||
register_complex,
|
||||
register_error,
|
||||
register_force_test,
|
||||
register_simple,
|
||||
split_complex_arg,
|
||||
split_complex_tensor,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
def register_binary_linear(op: OpType) -> Callable[..., Any]:
|
||||
def impl_with_alpha(
|
||||
lhs: ComplexTensor,
|
||||
rhs: ComplexTensor,
|
||||
*args: Any,
|
||||
alpha: int | float | complex,
|
||||
**kwargs: Any,
|
||||
) -> ComplexTensor:
|
||||
return op(lhs, aten.mul(rhs, alpha, *args, **kwargs), *args, **kwargs)
|
||||
|
||||
def impl(
|
||||
lhs: ComplexTensor, rhs: ComplexTensor, *args: Any, **kwargs: Any
|
||||
) -> ComplexTensor:
|
||||
alpha = kwargs.pop("alpha", None)
|
||||
if alpha is not None:
|
||||
return impl_with_alpha(lhs, rhs, *args, alpha=alpha, **kwargs)
|
||||
a_r, a_i = split_complex_arg(lhs)
|
||||
b_r, b_i = split_complex_arg(rhs)
|
||||
out_dt, (a_r, a_i, b_r, b_i) = promote_tensors(a_r, a_i, b_r, b_i)
|
||||
u = op(a_r, b_r, *args, **kwargs)
|
||||
v = op(a_i, b_i, *args, **kwargs)
|
||||
return ComplexTensor(u.to(out_dt), v.to(out_dt))
|
||||
|
||||
return register_complex(op, impl)
|
||||
|
||||
|
||||
@register_complex(aten.real)
|
||||
def real_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
re, _ = split_complex_tensor(self)
|
||||
return re
|
||||
|
||||
|
||||
@register_complex(aten.imag)
|
||||
def imag_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
_, im = split_complex_tensor(self)
|
||||
return im
|
||||
|
||||
|
||||
@register_complex(aten.is_pinned)
|
||||
def is_pinned_impl(self: ComplexTensor, device: torch.device | None = None) -> bool:
|
||||
return self.is_pinned(device)
|
||||
|
||||
|
||||
SIMPLE_OPS_LIST = [
|
||||
aten.slice,
|
||||
aten.flatten,
|
||||
aten.view,
|
||||
aten.diagonal,
|
||||
aten.expand,
|
||||
aten.unsqueeze,
|
||||
aten.unsqueeze_,
|
||||
aten.mean,
|
||||
aten.sum,
|
||||
aten.clone,
|
||||
aten.neg,
|
||||
aten.flip,
|
||||
aten.permute,
|
||||
aten.repeat,
|
||||
aten.index_select,
|
||||
aten.split,
|
||||
aten.split_with_sizes,
|
||||
aten.cumsum,
|
||||
aten.detach,
|
||||
aten.select,
|
||||
aten.squeeze,
|
||||
aten.zero_,
|
||||
aten.transpose,
|
||||
aten.t,
|
||||
aten.gather,
|
||||
]
|
||||
|
||||
for simple_op in SIMPLE_OPS_LIST:
|
||||
globals()[_get_func_name(simple_op)] = register_simple(simple_op)
|
||||
|
||||
# TODO (hameerabbasi): Not being tested
|
||||
SIMPLE_FORCE_TESTED_OPS = [
|
||||
aten.copy,
|
||||
aten.col2im,
|
||||
aten.alias,
|
||||
aten.lift_fresh,
|
||||
aten._unsafe_view,
|
||||
aten.index,
|
||||
aten._neg_view,
|
||||
aten.avg_pool2d,
|
||||
aten.avg_pool3d,
|
||||
aten.avg_pool2d_backward,
|
||||
aten.avg_pool3d_backward,
|
||||
aten.masked_scatter_backward,
|
||||
aten.select_backward,
|
||||
aten.slice_backward,
|
||||
aten.embedding,
|
||||
]
|
||||
|
||||
for simple_op in SIMPLE_FORCE_TESTED_OPS:
|
||||
globals()[_get_func_name(simple_op)] = register_force_test(
|
||||
simple_op, register_simple(simple_op)
|
||||
)
|
||||
|
||||
del simple_op
|
||||
|
||||
# some binary ops which we can stamp out
|
||||
mul_impl = register_binary_nonlinear(aten.mul)
|
||||
mul__impl = register_binary_nonlinear(aten.mul_)
|
||||
mm_impl = register_binary_nonlinear(aten.mm)
|
||||
dot_impl = register_binary_nonlinear(aten.dot)
|
||||
bmm_impl = register_binary_nonlinear(aten.bmm)
|
||||
|
||||
# TODO (hameerabbasi): Not being tested
|
||||
convolution_impl = register_force_test(
|
||||
aten.convolution, register_binary_nonlinear(aten.convolution)
|
||||
)
|
||||
|
||||
slice_scatter_impl = register_force_test(
|
||||
aten.slice_scatter, register_binary_linear(aten.slice_scatter)
|
||||
)
|
||||
select_scatter_impl = register_force_test(
|
||||
aten.select_scatter, register_binary_linear(aten.select_scatter)
|
||||
)
|
||||
|
||||
add_impl = register_binary_linear(aten.add)
|
||||
add__impl = register_binary_linear(aten.add_)
|
||||
sub_impl = register_binary_linear(aten.sub)
|
||||
sub__impl = register_binary_linear(aten.sub_)
|
||||
diagonal_scatter_impl = register_binary_linear(aten.diagonal_scatter)
|
||||
fill__impl = register_binary_linear(aten.fill_)
|
||||
|
||||
|
||||
@register_complex(aten.rsub)
|
||||
def rsub_impl(
|
||||
lhs: ComplexTensor, rhs: ComplexTensor, alpha: int | float | complex | None = None
|
||||
) -> ComplexTensor:
|
||||
if alpha is None:
|
||||
return torch.sub(rhs, lhs) # type: ignore[bad-return]
|
||||
return torch.sub(rhs, lhs, alpha=alpha) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.div)
|
||||
@register_complex(aten.true_divide)
|
||||
def div_impl(
|
||||
lhs: ComplexTensor, rhs: ComplexTensor, *, rounding_mode: str | None = None
|
||||
) -> ComplexTensor:
|
||||
if rounding_mode is not None:
|
||||
raise NotImplementedError(
|
||||
"`rounding_mode` other than `None` not implemented for`ComplexTensor`."
|
||||
)
|
||||
a_r, a_i = split_complex_arg(lhs)
|
||||
if not is_complex(rhs):
|
||||
return ComplexTensor(a_r / rhs, a_i / rhs)
|
||||
b_r, b_i = split_complex_arg(rhs)
|
||||
out_dt, (a_r, a_i, b_r, b_i) = promote_tensors(a_r, a_i, b_r, b_i)
|
||||
num_r = a_r * b_r + a_i * b_i
|
||||
num_i = a_i * b_r - a_r * b_i
|
||||
den = b_r * b_r + b_i * b_i
|
||||
return ComplexTensor(
|
||||
(num_r / den).to(out_dt),
|
||||
(num_i / den).to(out_dt),
|
||||
)
|
||||
|
||||
|
||||
@register_complex(aten.reciprocal)
|
||||
def reciprocal_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
self_r, self_i = split_complex_tensor(self)
|
||||
out_dt, (self_r, self_i) = promote_tensors(self_r, self_i)
|
||||
den = self_r * self_r + self_i * self_i
|
||||
return ComplexTensor(
|
||||
aten.div(self_r, den).to(out_dt),
|
||||
aten.div(-self_i, den).to(out_dt),
|
||||
)
|
||||
|
||||
|
||||
# reductions
|
||||
@register_complex(aten.prod)
|
||||
def prod_impl(self: ComplexTensor, *args: Any, **kwargs: Any) -> ComplexTensor:
|
||||
out_dt, (self,) = promote_tensors(self)
|
||||
dtype = kwargs.pop("dtype", out_dt)
|
||||
kwargs["dtype"] = complex_to_real_dtype(self.dtype)
|
||||
|
||||
prod_r = torch.prod(torch.abs(self), *args, **kwargs)
|
||||
sum_phi = torch.sum(torch.angle(self), *args, **kwargs)
|
||||
u = prod_r * torch.cos(sum_phi)
|
||||
v = prod_r * torch.sin(sum_phi)
|
||||
return ComplexTensor(u, v).to(dtype) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.pow)
|
||||
def pow_impl(self: ComplexTensor, exponent: ComplexTensor) -> ComplexTensor:
|
||||
out_dt, (self, exponent) = promote_tensors(self, exponent)
|
||||
return torch.exp(exponent * torch.log(self)).to(out_dt) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.cumprod)
|
||||
def cumprod_impl(self: ComplexTensor, *args: Any, **kwargs: Any) -> ComplexTensor:
|
||||
dtype = kwargs.pop("dtype", self.dtype)
|
||||
kwargs["dtype"] = complex_to_real_dtype(dtype)
|
||||
|
||||
prod_r = torch.cumprod(torch.abs(self), *args, **kwargs)
|
||||
sum_phi = torch.cumsum(torch.angle(self), *args, **kwargs)
|
||||
u = prod_r * torch.cos(sum_phi)
|
||||
v = prod_r * torch.sin(sum_phi)
|
||||
return ComplexTensor(u, v)
|
||||
|
||||
|
||||
# unary funcs,
|
||||
# most of these are simple or require some kind of identity
|
||||
@register_complex(aten.abs)
|
||||
def abs_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
out_dt, (x, y) = promote_tensors(x, y)
|
||||
result = torch.hypot(x, y)
|
||||
return result.to(out_dt)
|
||||
|
||||
|
||||
@register_complex(aten.angle)
|
||||
def angle_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
return torch.atan2(y, x)
|
||||
|
||||
|
||||
@register_complex(aten.acos)
|
||||
def acos_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
_, y = split_complex_tensor(self)
|
||||
acosh_z = torch.acosh(self)
|
||||
if not isinstance(acosh_z, ComplexTensor):
|
||||
raise AssertionError(f"acosh_z must be a ComplexTensor, got {type(acosh_z)}")
|
||||
acosh_z_re, acosh_z_im = split_complex_tensor(acosh_z)
|
||||
sign_im = 2 * torch.signbit(y) - 1
|
||||
return ComplexTensor(torch.abs(acosh_z_im), sign_im * torch.abs(acosh_z_re))
|
||||
|
||||
|
||||
@register_complex(aten.asin)
|
||||
def asin_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
asinh_iz = torch.asinh(ComplexTensor(-y, x))
|
||||
if not isinstance(asinh_iz, ComplexTensor):
|
||||
raise AssertionError(f"asinh_iz must be a ComplexTensor, got {type(asinh_iz)}")
|
||||
asinh_iz_re, asinh_iz_im = split_complex_tensor(asinh_iz)
|
||||
return ComplexTensor(asinh_iz_im, -asinh_iz_re)
|
||||
|
||||
|
||||
@register_complex(aten.atan)
|
||||
def atan_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
tanh_iz = torch.atanh(ComplexTensor(-y, x))
|
||||
if not isinstance(tanh_iz, ComplexTensor):
|
||||
raise AssertionError(f"tanh_iz must be a ComplexTensor, got {type(tanh_iz)}")
|
||||
tanh_iz_re, tanh_iz_im = split_complex_tensor(tanh_iz)
|
||||
return ComplexTensor(tanh_iz_im, -tanh_iz_re)
|
||||
|
||||
|
||||
@register_complex(aten.asinh)
|
||||
def asinh_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
out_dt, (self,) = promote_tensors(self)
|
||||
return torch.log(self + torch.sqrt(self * self + 1)).to(out_dt) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.acosh)
|
||||
def acosh_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
out_dt, (self,) = promote_tensors(self)
|
||||
return torch.log(self + torch.sqrt(self * self - 1)).to(out_dt) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.atanh)
|
||||
def atanh_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
out_dt, (x, y) = promote_tensors(x, y)
|
||||
|
||||
ret = 0.5 * (
|
||||
torch.log(ComplexTensor(1 + x, y)) - torch.log(ComplexTensor(1 - x, -y))
|
||||
)
|
||||
if not isinstance(ret, ComplexTensor):
|
||||
raise AssertionError(f"ret must be a ComplexTensor, got {type(ret)}")
|
||||
ret_re, ret_im = split_complex_tensor(ret)
|
||||
|
||||
return ComplexTensor(ret_re.to(out_dt), ret_im.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.cos)
|
||||
def cos_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
return torch.cosh(ComplexTensor(-y, x)) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.cosh)
|
||||
def cosh_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
out_dt, (x, y) = promote_tensors(x, y)
|
||||
u = torch.cosh(x) * torch.cos(y)
|
||||
v = torch.sinh(x) * torch.sin(y)
|
||||
return ComplexTensor(u.to(out_dt), v.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.sin)
|
||||
def sin_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
sinh_iz = torch.sinh(ComplexTensor(-y, x))
|
||||
if not isinstance(sinh_iz, ComplexTensor):
|
||||
raise AssertionError(f"sinh_iz must be a ComplexTensor, got {type(sinh_iz)}")
|
||||
sinh_iz_re, sinh_iz_im = split_complex_tensor(sinh_iz)
|
||||
return ComplexTensor(sinh_iz_im, -sinh_iz_re)
|
||||
|
||||
|
||||
@register_complex(aten.sinh)
|
||||
def sinh_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
out_dt, (x, y) = promote_tensors(x, y)
|
||||
u = torch.sinh(x) * torch.cos(y)
|
||||
v = torch.cosh(x) * torch.sin(y)
|
||||
return ComplexTensor(u.to(out_dt), v.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.tan)
|
||||
def tan_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
tanh_iz = torch.tanh(ComplexTensor(-y, x))
|
||||
if not isinstance(tanh_iz, ComplexTensor):
|
||||
raise AssertionError(f"tanh_iz must be a ComplexTensor, got {type(tanh_iz)}")
|
||||
tanh_iz_re, tanh_iz_im = split_complex_tensor(tanh_iz)
|
||||
return ComplexTensor(tanh_iz_im, -tanh_iz_re)
|
||||
|
||||
|
||||
@register_complex(aten.tanh)
|
||||
def tanh_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
out_dt, (x, y) = promote_tensors(x, y)
|
||||
|
||||
_2x = 2 * x
|
||||
_2y = 2 * y
|
||||
_d = torch.cosh(_2x) + torch.cos(_2y)
|
||||
_2xsh = torch.sinh(_2x)
|
||||
|
||||
out_re = _2xsh / _d
|
||||
out_im = torch.sin(_2y) / _d
|
||||
|
||||
return ComplexTensor(out_re.to(out_dt), out_im.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.exp)
|
||||
def exp_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
out_dt, (x, y) = promote_tensors(x, y)
|
||||
ex = torch.exp(x)
|
||||
u = ex * torch.cos(y)
|
||||
v = ex * torch.sin(y)
|
||||
return ComplexTensor(u.to(out_dt), v.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.expm1)
|
||||
def expm1_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
out_dt, (x, y) = promote_tensors(x, y)
|
||||
# TODO (hameerabbasi): The two lines below may have numerical issues
|
||||
ex = torch.exp(x)
|
||||
u = ex * torch.cos(y) - 1
|
||||
v = ex * torch.sin(y)
|
||||
return ComplexTensor(u.to(out_dt), v.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.log)
|
||||
def log_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
out_dt, (self,) = promote_tensors(self)
|
||||
re = torch.log(torch.abs(self))
|
||||
im = torch.angle(self)
|
||||
return ComplexTensor(re, im).to(out_dt) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.log1p)
|
||||
def log1p_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
# TODO (hameerabbasi): The line below may have numerical issues
|
||||
return torch.log(ComplexTensor(x + 1, y)) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.any)
|
||||
def any_impl(self: ComplexTensor, *args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
return torch.any(x, *args, **kwargs) | torch.any(y, *args, **kwargs)
|
||||
|
||||
|
||||
@register_complex(aten.all)
|
||||
def all_impl(self: ComplexTensor, *args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
return torch.any(x, *args, **kwargs) & torch.any(y, *args, **kwargs)
|
||||
|
||||
|
||||
@register_complex(aten.eq)
|
||||
def eq_impl(
|
||||
self: ComplexTensor, rhs: ComplexTensor, *args: Any, **kwargs: Any
|
||||
) -> torch.Tensor:
|
||||
a_r, a_i = split_complex_arg(self)
|
||||
b_r, b_i = split_complex_arg(rhs)
|
||||
return torch.eq(a_r, b_r, *args, **kwargs) & torch.eq(a_i, b_i, *args, **kwargs)
|
||||
|
||||
|
||||
@register_complex(aten.ne)
|
||||
def ne_impl(
|
||||
self: ComplexTensor, rhs: ComplexTensor, *args: Any, **kwargs: Any
|
||||
) -> torch.Tensor:
|
||||
a_r, a_i = split_complex_tensor(self)
|
||||
b_r, b_i = split_complex_arg(rhs)
|
||||
return torch.ne(a_r, b_r, *args, **kwargs) | torch.ne(a_i, b_i, *args, **kwargs)
|
||||
|
||||
|
||||
@register_complex(aten.isnan)
|
||||
def isnan_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
re, im = split_complex_tensor(self)
|
||||
return torch.isnan(re) | torch.isnan(im)
|
||||
|
||||
|
||||
@register_complex(aten.isinf)
|
||||
def isinf_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
re, im = split_complex_tensor(self)
|
||||
return torch.isinf(re) | torch.isinf(im)
|
||||
|
||||
|
||||
@register_complex(aten.isfinite)
|
||||
def isfinite_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
re, im = split_complex_tensor(self)
|
||||
return torch.isfinite(re) & torch.isfinite(im)
|
||||
|
||||
|
||||
@register_complex(aten.isclose)
|
||||
def isclose_impl(
|
||||
self: ComplexTensor,
|
||||
rhs: ComplexTensor,
|
||||
rtol: float = 1e-5,
|
||||
atol: float = 1e-8,
|
||||
equal_nan: bool = False,
|
||||
) -> torch.Tensor:
|
||||
abs_diff = torch.abs(self - rhs)
|
||||
abs_other = torch.abs(rhs)
|
||||
basic_condition = abs_diff <= (rtol * abs_other + atol)
|
||||
|
||||
# This is the nontrivial part
|
||||
if equal_nan:
|
||||
a_r, a_i = split_complex_tensor(self)
|
||||
b_r, b_i = split_complex_arg(rhs)
|
||||
|
||||
a_r_nan = torch.isnan(a_r)
|
||||
b_r_nan = torch.isnan(b_r)
|
||||
a_i_nan = torch.isnan(a_i)
|
||||
b_i_nan = torch.isnan(b_i)
|
||||
a_nan = a_r_nan | a_i_nan
|
||||
|
||||
# This logical expression makes sure that the isnan of both the real and imaginary parts
|
||||
# matches (so 1 + nan*i doesn't equal nan + 1*i)
|
||||
equal_nan_condition = ((a_r_nan == b_r_nan) & (a_i_nan == b_i_nan)) & a_nan
|
||||
return basic_condition | equal_nan_condition
|
||||
|
||||
return basic_condition
|
||||
|
||||
|
||||
ERROR_OPS_LIST = [
|
||||
aten.lt,
|
||||
aten.le,
|
||||
aten.gt,
|
||||
aten.ge,
|
||||
aten.amin,
|
||||
aten.amax,
|
||||
aten.clamp,
|
||||
aten.ceil,
|
||||
aten.floor,
|
||||
aten.minimum,
|
||||
aten.maximum,
|
||||
aten.trunc,
|
||||
aten.sign,
|
||||
aten.argmax,
|
||||
aten.argmin,
|
||||
aten.sort,
|
||||
aten.topk,
|
||||
aten.round,
|
||||
aten.fmod,
|
||||
]
|
||||
|
||||
|
||||
ERROR_TYPES = {
|
||||
aten.minimum: RuntimeError,
|
||||
aten.maximum: RuntimeError,
|
||||
aten.argmax: RuntimeError,
|
||||
aten.argmin: RuntimeError,
|
||||
aten.sort: RuntimeError,
|
||||
aten.topk: RuntimeError,
|
||||
}
|
||||
|
||||
|
||||
for err_op in ERROR_OPS_LIST:
|
||||
globals()[_get_func_name(err_op)] = register_error(
|
||||
err_op, ERROR_TYPES.get(err_op, NotImplementedError)
|
||||
)
|
||||
|
||||
del err_op
|
||||
|
||||
|
||||
@register_complex(aten.masked_scatter)
|
||||
def masked_scatter_impl(
|
||||
self: ComplexTensor, mask: torch.Tensor, source: ComplexTensor
|
||||
) -> ComplexTensor:
|
||||
self_r, self_i = split_complex_tensor(self)
|
||||
source_r, source_i = split_complex_arg(source)
|
||||
ret_r = torch.masked_scatter(self_r, mask, source_r)
|
||||
ret_i = torch.masked_scatter(self_i, mask, source_i)
|
||||
|
||||
return ComplexTensor(ret_r, ret_i)
|
||||
|
||||
|
||||
@register_complex(aten.where)
|
||||
def where_impl(mask: torch.Tensor, x: ComplexTensor, y: ComplexTensor) -> ComplexTensor:
|
||||
x_r, x_i = split_complex_arg(x)
|
||||
y_r, y_i = split_complex_arg(y)
|
||||
|
||||
ret_r = torch.where(mask, x_r, y_r)
|
||||
ret_i = torch.where(mask, x_i, y_i)
|
||||
|
||||
return ComplexTensor(ret_r, ret_i)
|
||||
|
||||
|
||||
@register_complex(aten.full_like)
|
||||
def full_like_impl(
|
||||
input: ComplexTensor,
|
||||
fill_value: complex,
|
||||
*args: Any,
|
||||
dtype: torch.dtype | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor | ComplexTensor:
|
||||
# Note: Cannot be merged with the cases below due to the `fill_value` argument
|
||||
input_r, input_i = split_complex_tensor(input)
|
||||
if dtype is not None and dtype not in COMPLEX_TO_REAL:
|
||||
return torch.full_like(input_r, fill_value, *args, dtype=dtype, **kwargs)
|
||||
|
||||
if dtype is not None:
|
||||
kwargs["dtype"] = COMPLEX_TO_REAL[dtype]
|
||||
|
||||
fv_r, fv_i = split_complex_arg(fill_value)
|
||||
ret_r = torch.full_like(input_r, fv_r, *args, **kwargs)
|
||||
ret_i = torch.full_like(input_i, fv_i, *args, **kwargs)
|
||||
|
||||
return ComplexTensor(ret_r, ret_i)
|
||||
|
||||
|
||||
def register_like(op: OpType) -> Callable[..., Any]:
|
||||
def impl(
|
||||
self: ComplexTensor, *args: Any, dtype: torch.dtype | None = None, **kwargs: Any
|
||||
) -> torch.Tensor | ComplexTensor:
|
||||
self_re, self_im = split_complex_tensor(self)
|
||||
|
||||
if dtype is not None and dtype not in COMPLEX_TO_REAL:
|
||||
return op(self_re, *args, dtype=dtype, **kwargs)
|
||||
|
||||
if dtype is not None:
|
||||
kwargs["dtype"] = COMPLEX_TO_REAL[dtype]
|
||||
|
||||
ret_re = op(self_re, *args, **kwargs)
|
||||
ret_im = op(self_im, *args, **kwargs)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
func_name = _get_func_name(op)
|
||||
impl.__name__ = func_name
|
||||
impl.__qualname__ = func_name
|
||||
|
||||
return register_complex(op, impl)
|
||||
|
||||
|
||||
LIKE_OPS_LIST = [
|
||||
aten.empty_like,
|
||||
aten.zeros_like,
|
||||
aten.randn_like,
|
||||
aten.new_zeros,
|
||||
]
|
||||
|
||||
for like_op in LIKE_OPS_LIST:
|
||||
globals()[_get_func_name(like_op)] = register_like(like_op)
|
||||
|
||||
del like_op
|
||||
|
||||
|
||||
@register_complex(aten.cat)
|
||||
def cat_impl(tensors: Sequence[ComplexTensor], dim: int = 0) -> ComplexTensor:
|
||||
tensors_r = []
|
||||
tensors_i = []
|
||||
|
||||
for t in tensors:
|
||||
t_r, t_i = split_complex_arg(t)
|
||||
tensors_r.append(t_r)
|
||||
tensors_i.append(t_i)
|
||||
|
||||
ret_r = torch.cat(tensors_r, dim=dim)
|
||||
ret_i = torch.cat(tensors_i, dim=dim)
|
||||
|
||||
return ComplexTensor(ret_r, ret_i)
|
||||
|
||||
|
||||
@register_complex(aten.sgn)
|
||||
def sgn_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
self_r, self_i = split_complex_tensor(self)
|
||||
out_dt, (self_r, self_i) = promote_tensors(self_r, self_i)
|
||||
abs_self = torch.abs(ComplexTensor(self_r, self_i))
|
||||
mask = (self_r != 0) | (self_i != 0)
|
||||
masked_sgn = ComplexTensor(
|
||||
(self_r / abs_self).to(out_dt), (self_i / abs_self).to(out_dt)
|
||||
)
|
||||
return torch.where(mask, masked_sgn, 0) # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.sqrt)
|
||||
def sqrt_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
self_r, self_i = split_complex_tensor(self)
|
||||
out_dt, (self_r, self_i) = promote_tensors(self_r, self_i)
|
||||
self = ComplexTensor(self_r, self_i)
|
||||
self_abs_sqrt = torch.sqrt(torch.abs(self))
|
||||
self_half_angle = 0.5 * torch.angle(self)
|
||||
|
||||
ret_r = self_abs_sqrt * torch.cos(self_half_angle)
|
||||
ret_i = self_abs_sqrt * torch.sin(self_half_angle)
|
||||
|
||||
return ComplexTensor(ret_r.to(out_dt), ret_i.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.rsqrt)
|
||||
def rsqrt_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
self_r, self_i = split_complex_tensor(self)
|
||||
out_dt, (self_r, self_i) = promote_tensors(self_r, self_i)
|
||||
self = ComplexTensor(self_r, self_i)
|
||||
self_abs_rsqrt = torch.rsqrt(torch.abs(self))
|
||||
self_neg_half_angle = -0.5 * torch.angle(self)
|
||||
|
||||
ret_r = self_abs_rsqrt * torch.cos(self_neg_half_angle)
|
||||
ret_i = self_abs_rsqrt * torch.sin(self_neg_half_angle)
|
||||
|
||||
return ComplexTensor(ret_r.to(out_dt), ret_i.to(out_dt))
|
||||
|
||||
|
||||
@register_complex(aten.addmm)
|
||||
def addmm_impl(
|
||||
input: ComplexTensor,
|
||||
mat1: ComplexTensor,
|
||||
mat2: ComplexTensor,
|
||||
out_dtype: torch.dtype | None = None,
|
||||
beta: int | float | complex = 1,
|
||||
alpha: int | float | complex = 1,
|
||||
) -> ComplexTensor:
|
||||
ret = beta * input + alpha * torch.mm(mat1, mat2)
|
||||
if not isinstance(ret, ComplexTensor):
|
||||
raise AssertionError(f"ret must be a ComplexTensor, got {type(ret)}")
|
||||
ret_r, ret_i = split_complex_tensor(ret)
|
||||
if out_dtype is not None:
|
||||
out_dtype = COMPLEX_TO_REAL[out_dtype]
|
||||
ret_r, ret_i = ret_r.to(out_dtype), ret_i.to(out_dtype)
|
||||
return ComplexTensor(ret_r, ret_i)
|
||||
|
||||
|
||||
def elemwise_nonzero(self: ComplexTensor) -> torch.Tensor:
|
||||
re, im = split_complex_tensor(self)
|
||||
return (re != 0) | (im != 0)
|
||||
|
||||
|
||||
def register_nonzero_impl(op: OpType) -> Callable[..., Any]:
|
||||
def nonzero_impl(
|
||||
self: ComplexTensor, other: ComplexTensor, *args: Any, **kwargs: Any
|
||||
) -> torch.Tensor:
|
||||
return op(elemwise_nonzero(self), elemwise_nonzero(other), *args, **kwargs)
|
||||
|
||||
func_name = _get_func_name(op)
|
||||
nonzero_impl.__name__ = func_name
|
||||
nonzero_impl.__qualname__ = func_name
|
||||
|
||||
return register_complex(op, nonzero_impl)
|
||||
|
||||
|
||||
logical_and_impl = register_nonzero_impl(aten.logical_and)
|
||||
logical_or_impl = register_nonzero_impl(aten.logical_or)
|
||||
logical_xor_impl = register_nonzero_impl(aten.logical_xor)
|
||||
|
||||
|
||||
@register_complex(aten.logical_not)
|
||||
def logical_not_impl(self: ComplexTensor, *args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
return torch.logical_not(elemwise_nonzero(self), *args, **kwargs)
|
||||
|
||||
|
||||
@register_complex(aten.view_as_real)
|
||||
def view_as_real_impl(self: ComplexTensor) -> torch.Tensor:
|
||||
re, im = split_complex_tensor(self)
|
||||
return torch.stack([re, im], dim=-1)
|
||||
|
||||
|
||||
@register_complex(aten.linalg_vector_norm)
|
||||
def linalg_vector_norm_impl(
|
||||
self: ComplexTensor, *args: Any, **kwargs: Any
|
||||
) -> torch.Tensor:
|
||||
return torch.linalg.vector_norm(torch.abs(self), *args, **kwargs)
|
||||
|
||||
|
||||
@register_force_test(aten.copy_)
|
||||
def copy__impl(
|
||||
self: ComplexTensor | torch.Tensor,
|
||||
src: ComplexTensor | torch.Tensor,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> ComplexTensor | torch.Tensor:
|
||||
if not self.dtype.is_complex:
|
||||
warnings.warn(
|
||||
"Casting complex values to real discards the imaginary part", UserWarning
|
||||
)
|
||||
src_re, src_im = split_complex_arg(src)
|
||||
return self.copy_(src_re)
|
||||
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
src_re, src_im = split_complex_arg(src)
|
||||
|
||||
ret_re = self_re.copy_(src_re, *args, **kwargs)
|
||||
ret_im = self_im.copy_(src_im, *args, **kwargs)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
|
||||
@register_complex(aten._local_scalar_dense)
|
||||
def _local_scalar_dense_impl(self: ComplexTensor, *args: Any, **kwargs: Any) -> complex:
|
||||
x, y = split_complex_tensor(self)
|
||||
u = aten._local_scalar_dense(x, *args, **kwargs)
|
||||
v = aten._local_scalar_dense(y, *args, **kwargs)
|
||||
return complex(u, v)
|
||||
|
||||
|
||||
@register_complex(aten.allclose)
|
||||
def allclose_impl(
|
||||
input: torch.Tensor,
|
||||
other: torch.Tensor,
|
||||
rtol: float = 1e-05,
|
||||
atol: float = 1e-08,
|
||||
equal_nan: bool = False,
|
||||
) -> bool:
|
||||
# pyrefly: ignore [bad-return]
|
||||
return torch.all(
|
||||
torch.isclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan)
|
||||
).item() # type: ignore[bad-return]
|
||||
|
||||
|
||||
@register_complex(aten.stack)
|
||||
def stack_impl(self: list[ComplexTensor], *args: Any, **kwargs: Any) -> ComplexTensor:
|
||||
re_im_tuples = [split_complex_arg(self_i) for self_i in self]
|
||||
u = torch.stack([c[0] for c in re_im_tuples], *args, **kwargs)
|
||||
v = torch.stack([c[1] for c in re_im_tuples], *args, **kwargs)
|
||||
return ComplexTensor(u, v)
|
||||
|
||||
|
||||
# TODO (hameerabbasi): Not being tested
|
||||
@register_complex(aten._conj_physical)
|
||||
@register_complex(aten.conj_physical)
|
||||
def conj_physical_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
re, im = split_complex_tensor(self)
|
||||
return ComplexTensor(re, -im)
|
||||
|
||||
|
||||
# TODO (hameerabbasi): Not being tested
|
||||
@register_complex(aten._conj)
|
||||
def _conj_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
re, im = split_complex_tensor(self)
|
||||
return ComplexTensor(re, torch._neg_view(im))
|
||||
|
||||
|
||||
@register_complex(aten.index_add)
|
||||
def index_add_impl(
|
||||
self: ComplexTensor,
|
||||
dim: int,
|
||||
index: torch.Tensor,
|
||||
source: ComplexTensor,
|
||||
**kwargs: Any,
|
||||
) -> ComplexTensor:
|
||||
alpha = kwargs.pop("alpha", None)
|
||||
if alpha is not None:
|
||||
source = source * alpha
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
source_re, source_im = split_complex_arg(source)
|
||||
|
||||
ret_re = self_re.index_add(dim, index, source_re)
|
||||
ret_im = self_im.index_add(dim, index, source_im)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
|
||||
# TODO (hameerabbasi): Not being tested
|
||||
@register_complex(aten.index_add_)
|
||||
def index_add__impl(
|
||||
self: ComplexTensor,
|
||||
dim: int,
|
||||
index: torch.Tensor,
|
||||
source: ComplexTensor,
|
||||
**kwargs: Any,
|
||||
) -> ComplexTensor:
|
||||
alpha = kwargs.pop("alpha", None)
|
||||
if alpha is not None:
|
||||
source = source * alpha
|
||||
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
source_re, source_im = split_complex_arg(source)
|
||||
|
||||
ret_re = self_re.index_add_(dim, index, source_re)
|
||||
ret_im = self_im.index_add_(dim, index, source_im)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
|
||||
@register_complex(aten.masked_fill)
|
||||
def masked_fill_impl(
|
||||
self: ComplexTensor, mask: torch.Tensor, value: complex
|
||||
) -> ComplexTensor:
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
value_re, value_im = split_complex_arg(value)
|
||||
|
||||
ret_re = self_re.masked_fill(mask, value_re)
|
||||
ret_im = self_im.masked_fill(mask, value_im)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
|
||||
# TODO (hameerabbasi): Not being tested
|
||||
@register_complex(aten.masked_fill_)
|
||||
def masked_fill__impl(
|
||||
self: ComplexTensor, mask: torch.Tensor, value: complex
|
||||
) -> ComplexTensor:
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
value_re, value_im = split_complex_arg(value)
|
||||
|
||||
ret_re = self_re.masked_fill_(mask, value_re)
|
||||
ret_im = self_im.masked_fill_(mask, value_im)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
|
||||
@register_complex(aten.constant_pad_nd)
|
||||
def constant_pad_nd_impl(
|
||||
self: ComplexTensor, pad: Sequence[int], value: complex | None = None
|
||||
) -> ComplexTensor:
|
||||
self_re, self_im = split_complex_tensor(self)
|
||||
if value is None:
|
||||
ret_re = aten.constant_pad_nd(self_re, pad)
|
||||
ret_im = aten.constant_pad_nd(self_im, pad)
|
||||
else:
|
||||
value_re, value_im = split_complex_arg(value)
|
||||
ret_re = aten.constant_pad_nd(self_re, pad, value_re)
|
||||
ret_im = aten.constant_pad_nd(self_im, pad, value_im)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
|
||||
@register_complex(aten.var)
|
||||
def var_impl(self: ComplexTensor, *args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
self_re, self_im = split_complex_tensor(self)
|
||||
return torch.var(self_re, *args, **kwargs) + torch.var(self_im, *args, **kwargs)
|
||||
|
||||
|
||||
@register_complex(aten.scatter_add)
|
||||
def scatter_add_impl(
|
||||
self: ComplexTensor, dim: int, index: torch.Tensor, src: ComplexTensor
|
||||
) -> ComplexTensor:
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
src_re, src_im = split_complex_arg(src)
|
||||
|
||||
ret_re = torch.scatter_add(self_re, dim, index, src_re)
|
||||
ret_im = torch.scatter_add(self_im, dim, index, src_im)
|
||||
|
||||
return ComplexTensor(ret_re, ret_im)
|
||||
|
||||
|
||||
@register_complex(aten.scatter_add_)
|
||||
def scatter_add__impl(
|
||||
self: ComplexTensor, dim: int, index: torch.Tensor, src: ComplexTensor
|
||||
) -> ComplexTensor:
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
src_re, src_im = split_complex_arg(src)
|
||||
|
||||
out_re = self_re.scatter_add_(dim, index, src_re)
|
||||
out_im = self_im.scatter_add_(dim, index, src_im)
|
||||
|
||||
return ComplexTensor(out_re, out_im)
|
||||
|
||||
|
||||
@register_complex(aten.index_put_)
|
||||
def index_put__impl(
|
||||
self: ComplexTensor,
|
||||
indices: tuple[torch.Tensor, ...],
|
||||
values: ComplexTensor,
|
||||
accumulate: bool = False,
|
||||
) -> ComplexTensor:
|
||||
self_re, self_im = split_complex_arg(self)
|
||||
values_re, values_im = split_complex_arg(values)
|
||||
|
||||
out_re = self_re.index_put_(indices, values_re, accumulate=accumulate)
|
||||
out_im = self_im.index_put_(indices, values_im, accumulate=accumulate)
|
||||
|
||||
return ComplexTensor(out_re, out_im)
|
||||
|
||||
|
||||
@register_complex(aten.tanh_backward)
|
||||
def tanh_backward(out_grad: ComplexTensor, y: ComplexTensor) -> ComplexTensor:
|
||||
# pyrefly: ignore[bad-return]
|
||||
return out_grad * (1.0 - y * y).conj_physical()
|
||||
|
||||
|
||||
@register_complex(aten.diagonal_backward)
|
||||
def diagonal_backward(
|
||||
grad_output: torch.Tensor, input_sizes: list[int], offset: int, dim1: int, dim2: int
|
||||
) -> torch.Tensor:
|
||||
grad_input = grad_output.new_zeros(input_sizes)
|
||||
return torch.diagonal_scatter(grad_input, grad_output, offset, dim1, dim2)
|
||||
|
||||
|
||||
def _dt_to_real(dt: torch.dtype | Any) -> torch.dtype | Any:
|
||||
if not isinstance(dt, torch.dtype):
|
||||
return dt
|
||||
|
||||
return COMPLEX_TO_REAL[dt]
|
||||
|
||||
|
||||
def register_to_impl(op: OpType) -> Callable[..., Any]:
|
||||
"""Register an op similar to `aten.to`, but may have different signatures."""
|
||||
|
||||
def impl(
|
||||
self: ComplexTensor, *args: Any, **kwargs: Any
|
||||
) -> torch.Tensor | ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
try:
|
||||
args = tuple(_dt_to_real(a) for a in args)
|
||||
kwargs = {k: _dt_to_real(v) for k, v in kwargs.items()}
|
||||
except KeyError:
|
||||
return op(x, *args, **kwargs)
|
||||
|
||||
return ComplexTensor(op(x, *args, **kwargs), op(y, *args, **kwargs))
|
||||
|
||||
func_name = _get_func_name(op)
|
||||
impl.__name__ = func_name
|
||||
impl.__qualname__ = func_name
|
||||
|
||||
return register_complex(op, impl)
|
||||
|
||||
|
||||
to_impl = register_to_impl(aten.to)
|
||||
_to_copy_impl = register_to_impl(aten._to_copy)
|
||||
+335
@@ -0,0 +1,335 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any, overload, TypeAlias
|
||||
from typing_extensions import Never, ParamSpec, TypeIs, TypeVar
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch._decomp import get_decompositions
|
||||
from torch._ops import OpOverload, OpOverloadPacket
|
||||
from torch._refs import is_complex as _is_complex
|
||||
from torch.types import Number
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
|
||||
|
||||
from .._core import ComplexTensor
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_R = TypeVar("_R")
|
||||
|
||||
OpType: TypeAlias = OpOverloadPacket | OpOverload
|
||||
|
||||
# pyrefly: ignore [implicit-any]
|
||||
TableType: TypeAlias = dict[OpType, Callable]
|
||||
|
||||
# Mapping from ops to implementations
|
||||
COMPLEX_OPS_TABLE: TableType = {}
|
||||
|
||||
COMPLEX_TO_REAL = {
|
||||
torch.complex128: torch.float64,
|
||||
torch.complex64: torch.float32,
|
||||
torch.complex32: torch.float16,
|
||||
}
|
||||
|
||||
REAL_TO_COMPLEX = {v: k for k, v in COMPLEX_TO_REAL.items()}
|
||||
|
||||
# Used to promote dtypes in `promote_real_cpu_tensors`
|
||||
PROMOTE_TYPES = {
|
||||
torch.float16: torch.float32,
|
||||
torch.bfloat16: torch.float32,
|
||||
torch.complex32: torch.complex64,
|
||||
}
|
||||
|
||||
|
||||
def is_complex_tensor(obj: Any, /) -> TypeIs[ComplexTensor]:
|
||||
r"""Returns True if the input is a ComplexTensor, else False
|
||||
|
||||
Args:
|
||||
a: any input
|
||||
|
||||
Examples:
|
||||
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> from torch.complex import ComplexTensor
|
||||
>>> data = torch.zeros((3, 2), dtype=torch.complex64)
|
||||
>>> ct = ComplexTensor.from_interleaved(data)
|
||||
>>> is_complex_tensor(ct)
|
||||
True
|
||||
"""
|
||||
return isinstance(obj, ComplexTensor)
|
||||
|
||||
|
||||
@overload
|
||||
def promote_tensors(
|
||||
*tensors: ComplexTensor,
|
||||
) -> tuple[torch.dtype, tuple[ComplexTensor, ...]]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def promote_tensors(
|
||||
*tensors: Tensor,
|
||||
) -> tuple[torch.dtype, tuple[Tensor, ...]]: ...
|
||||
|
||||
|
||||
def promote_tensors(
|
||||
*tensors: Tensor | ComplexTensor,
|
||||
) -> tuple[torch.dtype, tuple[Tensor | ComplexTensor, ...]]:
|
||||
"""
|
||||
Promotes all tensors to a common dtype.
|
||||
Additionally promotes CPU tensors to at least `float32`.
|
||||
"""
|
||||
tensor = next(t for t in tensors if isinstance(t, Tensor))
|
||||
out_dt = tensor.dtype
|
||||
for t in tensors:
|
||||
if isinstance(t, Tensor):
|
||||
out_dt = torch.promote_types(out_dt, t.dtype)
|
||||
|
||||
prom_dt = PROMOTE_TYPES.get(out_dt, out_dt)
|
||||
return out_dt, tuple(
|
||||
t.to(prom_dt) if isinstance(t, Tensor) else torch.asarray(t, dtype=prom_dt)
|
||||
for t in tensors
|
||||
)
|
||||
|
||||
|
||||
def register_complex(
|
||||
op: OpType,
|
||||
func_impl: Callable[..., Any] | None = None,
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]] | Callable[..., Any]:
|
||||
"""Decorator to register an implementation for some ops in some dispatch tables"""
|
||||
|
||||
def inner(func: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
if COMPLEX_OPS_TABLE.get(op, func) is not func:
|
||||
raise RuntimeError(f"Attempted to register multiple functions for {op}")
|
||||
COMPLEX_OPS_TABLE[op] = func
|
||||
return func
|
||||
|
||||
if func_impl is None:
|
||||
return inner
|
||||
|
||||
return inner(func_impl)
|
||||
|
||||
|
||||
FORCE_TEST_LIST: list[OpType] = []
|
||||
|
||||
|
||||
def register_force_test(
|
||||
op: OpType, func_impl: Callable[..., Any] | None = None
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]] | Callable[..., Any]:
|
||||
"""Will attempt to test these ops even if they err on "normal" inputs"""
|
||||
FORCE_TEST_LIST.append(op)
|
||||
return register_complex(op, func_impl)
|
||||
|
||||
|
||||
DECOMPOSITIONS = get_decompositions(list(torch.ops.aten)) # type: ignore[no-matching-overload]
|
||||
|
||||
|
||||
def lookup_complex(
|
||||
func: OpOverload, *args: Any, **kwargs: Any
|
||||
) -> Callable[..., Any] | None:
|
||||
"""
|
||||
Lookup an impl from the table.
|
||||
|
||||
Try the particular overload first, then the overload packet.
|
||||
|
||||
If nothing is found, try the decompositions with both.
|
||||
"""
|
||||
return COMPLEX_OPS_TABLE.get(
|
||||
func,
|
||||
COMPLEX_OPS_TABLE.get(
|
||||
func.overloadpacket,
|
||||
DECOMPOSITIONS.get(func, DECOMPOSITIONS.get(func.overloadpacket)),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def is_complex(x: Any, /) -> bool:
|
||||
"""Utility to detect if a given object is (known) to be complex."""
|
||||
return (isinstance(x, Tensor) and _is_complex(x)) or isinstance(x, complex)
|
||||
|
||||
|
||||
@overload
|
||||
def split_complex_arg(
|
||||
arg: Tensor | ComplexTensor,
|
||||
) -> tuple[Tensor, Tensor]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def split_complex_arg(
|
||||
arg: complex | Number,
|
||||
) -> tuple[Number, Number]: ...
|
||||
|
||||
|
||||
def split_complex_arg(
|
||||
arg: Tensor | ComplexTensor | complex | Number,
|
||||
) -> tuple[Tensor, Tensor] | tuple[Number, Number]:
|
||||
"""
|
||||
Split a complex argument into a real/imaginary component.
|
||||
|
||||
If real, use zero for the imaginary part.
|
||||
"""
|
||||
if isinstance(arg, ComplexTensor):
|
||||
return split_complex_tensor(arg)
|
||||
if isinstance(arg, Tensor):
|
||||
if is_complex(arg):
|
||||
return arg.real, arg.imag
|
||||
return arg, torch.zeros_like(arg)
|
||||
# TODO (hameerabbasi): Should there be a `torch.SymComplex`?
|
||||
if isinstance(arg, complex):
|
||||
return arg.real, arg.imag
|
||||
if isinstance(arg, float | torch.SymFloat):
|
||||
return arg, 0.0
|
||||
if isinstance(arg, int | torch.SymInt):
|
||||
return arg, 0
|
||||
if isinstance(arg, bool | torch.SymBool):
|
||||
return arg, False
|
||||
raise TypeError(f"Expected tensor or number got, {type(arg)}")
|
||||
|
||||
|
||||
def split_complex_tensor(complex_tensor: ComplexTensor) -> tuple[Tensor, Tensor]:
|
||||
"""Split a ComplexTensor into its real and imaginary parts."""
|
||||
return complex_tensor.re, complex_tensor.im
|
||||
|
||||
|
||||
def complex_to_real_dtype(dtype: torch.dtype) -> torch.dtype:
|
||||
"""Convert a complex dtype to the dtype of its real part. Return other dtypes as-is."""
|
||||
return COMPLEX_TO_REAL.get(dtype, dtype)
|
||||
|
||||
|
||||
def _get_op_name(op: OpType) -> str:
|
||||
"""Get the op name from the op."""
|
||||
if isinstance(op, OpOverload):
|
||||
op = op.overloadpacket
|
||||
return str(op).split(".", 1)[1]
|
||||
|
||||
|
||||
def _get_func_name(op: OpType) -> str:
|
||||
"""Get the name of the implementation function from the op."""
|
||||
return f"{_get_op_name(op)}_impl"
|
||||
|
||||
|
||||
def register_error(
|
||||
op: OpType, exc_type: type[Exception] = NotImplementedError
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]] | Callable[..., Any]:
|
||||
msg = f"`aten.{_get_op_name(op)}` not implemented for `{ComplexTensor.__name__}`."
|
||||
|
||||
def ordered_impl(*args: Any, **kwargs: Any) -> Never:
|
||||
raise exc_type(msg)
|
||||
|
||||
func_name = _get_func_name(op)
|
||||
ordered_impl.__name__ = func_name
|
||||
ordered_impl.__qualname__ = func_name
|
||||
|
||||
return register_force_test(op, ordered_impl)
|
||||
|
||||
|
||||
def register_binary_nonlinear(
|
||||
op: OpType,
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]] | Callable[..., Any]:
|
||||
"""Register a "multiplication-style" op, e.g. aten.mul, aten.mm, ..."""
|
||||
|
||||
def impl(
|
||||
lhs: ComplexTensor, rhs: ComplexTensor, *args: Any, **kwargs: Any
|
||||
) -> ComplexTensor:
|
||||
a_r, a_i = split_complex_arg(lhs)
|
||||
b_r, b_i = split_complex_arg(rhs)
|
||||
out_dt, (a_r, a_i, b_r, b_i) = promote_tensors(a_r, a_i, b_r, b_i)
|
||||
real = op(a_r, b_r, *args, **kwargs) - op(a_i, b_i, *args, **kwargs)
|
||||
imag = op(a_r, b_i, *args, **kwargs) + op(a_i, b_r, *args, **kwargs)
|
||||
return ComplexTensor(real.to(out_dt), imag.to(out_dt))
|
||||
|
||||
func_name = _get_func_name(op)
|
||||
impl.__name__ = func_name
|
||||
impl.__qualname__ = func_name
|
||||
|
||||
return register_complex(op, impl)
|
||||
|
||||
|
||||
def register_simple(
|
||||
op: OpType,
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]] | Callable[..., Any]:
|
||||
"""Register an op which can be applied independently to the real and complex parts to get the result."""
|
||||
|
||||
def impl(
|
||||
self: ComplexTensor, *args: Any, dtype: torch.dtype | None = None, **kwargs: Any
|
||||
) -> ComplexTensor:
|
||||
x, y = split_complex_tensor(self)
|
||||
if dtype is not None and dtype not in COMPLEX_TO_REAL:
|
||||
raise RuntimeError(
|
||||
"Non-complex `dtype` specified, please write custom impl."
|
||||
)
|
||||
|
||||
if dtype in COMPLEX_TO_REAL:
|
||||
if dtype is None:
|
||||
raise AssertionError("dtype must not be None when in COMPLEX_TO_REAL")
|
||||
kwargs["dtype"] = COMPLEX_TO_REAL[dtype]
|
||||
|
||||
u = op(x, *args, **kwargs)
|
||||
v = op(y, *args, **kwargs)
|
||||
|
||||
u_flat, u_spec = tree_flatten(u)
|
||||
v_flat, v_spec = tree_flatten(v)
|
||||
if u_spec != v_spec:
|
||||
raise AssertionError(f"Tree specs must match: {u_spec} != {v_spec}")
|
||||
out_flat = [
|
||||
ComplexTensor(ui, vi) for ui, vi in zip(u_flat, v_flat, strict=False)
|
||||
]
|
||||
return tree_unflatten(out_flat, u_spec)
|
||||
|
||||
func_name = _get_func_name(op)
|
||||
impl.__name__ = func_name
|
||||
impl.__qualname__ = func_name
|
||||
|
||||
return register_complex(op, impl)
|
||||
|
||||
|
||||
def _as_complex_tensor(arg: Tensor | Any) -> Tensor | ComplexTensor | Any:
|
||||
"""Convert a Tensor with complex dtypes to a ComplexTensor. Pass along other args as-is."""
|
||||
if (
|
||||
not isinstance(arg, ComplexTensor)
|
||||
and isinstance(arg, Tensor)
|
||||
and arg.dtype in COMPLEX_TO_REAL
|
||||
):
|
||||
return ComplexTensor.from_interleaved(arg)
|
||||
return arg
|
||||
|
||||
|
||||
def _as_interleaved(arg: ComplexTensor | Any) -> Tensor | Any:
|
||||
"""Convert a ComplexTensor to a Tensor with a complex dtype. Pass other arguments as-is."""
|
||||
if isinstance(arg, ComplexTensor):
|
||||
return arg.as_interleaved()
|
||||
return arg
|
||||
|
||||
|
||||
class ComplexTensorMode(TorchDispatchMode):
|
||||
_compile: bool
|
||||
|
||||
""" A TorchDispatchMode to replace any Tensor that has a complex dtype with a ComplexTensor for the computation. """
|
||||
|
||||
def __init__(self, _dispatch_key: Any = None, *, _compile: bool = False) -> None:
|
||||
"""Initialize a ComplexTensorMode.
|
||||
|
||||
Args:
|
||||
_dispatch_key: passed on to TorchDispatchMode
|
||||
_compile: Compile the op before the computation
|
||||
"""
|
||||
super().__init__(_dispatch_key)
|
||||
self._compile = _compile
|
||||
|
||||
def __torch_dispatch__(
|
||||
self,
|
||||
func: OpOverload,
|
||||
types: tuple[type, ...],
|
||||
args: tuple[Any, ...] = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
# TODO (hameerabbasi): Test perf with `_compile` set to `True`
|
||||
if self._compile:
|
||||
func = torch.compile(func) # type: ignore[bad-assignment]
|
||||
|
||||
args = tree_map(_as_complex_tensor, args)
|
||||
kwargs = tree_map(_as_complex_tensor, kwargs)
|
||||
|
||||
return tree_map(_as_interleaved, func(*args, **kwargs))
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
import torch
|
||||
|
||||
from .._core import ComplexTensor
|
||||
from .common import (
|
||||
complex_to_real_dtype,
|
||||
register_complex,
|
||||
register_force_test,
|
||||
split_complex_tensor,
|
||||
)
|
||||
|
||||
|
||||
prims = torch.ops.prims
|
||||
aten = torch.ops.aten
|
||||
|
||||
|
||||
# TODO (hameerabbasi): Not being tested
|
||||
@register_force_test(prims.convert_element_type)
|
||||
def convert_element_type_impl(x: ComplexTensor, dtype: torch.dtype) -> ComplexTensor:
|
||||
dtype = complex_to_real_dtype(dtype)
|
||||
u, v = split_complex_tensor(x)
|
||||
u_out = prims.convert_element_type(u, dtype)
|
||||
v_out = prims.convert_element_type(v, dtype)
|
||||
|
||||
return ComplexTensor(u_out, v_out)
|
||||
|
||||
|
||||
@register_complex(prims.conj_physical)
|
||||
def conj_physical_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
return aten._conj_physical(self)
|
||||
|
||||
|
||||
@register_complex(prims.conj)
|
||||
def conj_impl(self: ComplexTensor) -> ComplexTensor:
|
||||
return aten._conj(self)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,326 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import warnings
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._subclasses.fake_tensor import (
|
||||
FakeTensor,
|
||||
FakeTensorMode,
|
||||
MetadataMismatchError,
|
||||
tree_flatten_only,
|
||||
UnsupportedFakeTensorException,
|
||||
)
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
|
||||
from torch._ops import OpOverload
|
||||
from torch.utils._pytree import PyTree
|
||||
|
||||
|
||||
aten = torch._ops.ops.aten
|
||||
|
||||
|
||||
def outputs_alias_inputs(outputs: PyTree, inputs: PyTree) -> bool:
|
||||
input_storages = {
|
||||
inp._typed_storage()._cdata
|
||||
for inp in tree_flatten_only(torch.Tensor, inputs)
|
||||
if torch._C._has_storage(inp)
|
||||
}
|
||||
return any(
|
||||
torch._C._has_storage(out) and out._typed_storage()._cdata in input_storages
|
||||
for out in tree_flatten_only(torch.Tensor, outputs)
|
||||
)
|
||||
|
||||
|
||||
def outputs_are_inputs(outputs: PyTree, inputs: PyTree) -> bool:
|
||||
input_ids = {id(inp) for inp in tree_flatten_only(torch.Tensor, inputs)}
|
||||
return any(id(out) in input_ids for out in tree_flatten_only(torch.Tensor, outputs))
|
||||
|
||||
|
||||
def output_alias_each_other(outputs: PyTree) -> bool:
|
||||
storages = set()
|
||||
for out in tree_flatten_only(torch.Tensor, outputs):
|
||||
if not torch._C._has_storage(out):
|
||||
continue
|
||||
stor = out._typed_storage()._cdata
|
||||
if stor in storages:
|
||||
return True
|
||||
storages.add(stor)
|
||||
return False
|
||||
|
||||
|
||||
def _check_alias_info(
|
||||
context: str,
|
||||
real_out: PyTree,
|
||||
real_in: PyTree,
|
||||
fake_out: PyTree,
|
||||
fake_in: PyTree,
|
||||
) -> None:
|
||||
r_aliasing = outputs_alias_inputs(real_out, real_in)
|
||||
f_aliasing = outputs_alias_inputs(fake_out, fake_in)
|
||||
if r_aliasing != f_aliasing:
|
||||
raise MetadataMismatchError(
|
||||
f"{context} mismatch in outputs_alias_inputs check {f_aliasing} != {r_aliasing}"
|
||||
)
|
||||
|
||||
r_identity_eq = outputs_are_inputs(real_out, real_in)
|
||||
f_identity_eq = outputs_are_inputs(fake_out, fake_in)
|
||||
if r_identity_eq != f_identity_eq:
|
||||
raise MetadataMismatchError(
|
||||
f"{context} mismatch in outputs_are_inputs check {f_identity_eq} != {r_identity_eq}"
|
||||
)
|
||||
|
||||
r_output_alias_each_other = output_alias_each_other(real_out)
|
||||
f_output_alias_each_other = output_alias_each_other(fake_out)
|
||||
if r_output_alias_each_other != f_output_alias_each_other:
|
||||
raise MetadataMismatchError(
|
||||
f"{context} mismatch in outputs_alias_each_other check "
|
||||
f"{f_output_alias_each_other} != {r_output_alias_each_other}"
|
||||
)
|
||||
|
||||
|
||||
def is_sdpa_error(func: OpOverload, idx: int, e: Exception) -> bool:
|
||||
if (
|
||||
(
|
||||
func is aten._scaled_dot_product_flash_attention.default
|
||||
or func is aten._flash_attention_forward.default
|
||||
)
|
||||
and idx in (6, 7)
|
||||
and "Devices" in repr(e)
|
||||
):
|
||||
return True
|
||||
if (
|
||||
(
|
||||
func is aten._scaled_dot_product_efficient_attention.default
|
||||
or func is aten._efficient_attention_forward.default
|
||||
)
|
||||
and idx in (2, 3)
|
||||
and "Devices" in repr(e)
|
||||
):
|
||||
return True
|
||||
if (
|
||||
func is aten._scaled_dot_product_cudnn_attention.default
|
||||
and idx in (6, 7)
|
||||
and "Devices" in repr(e)
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def try_convert_fake_to_real(
|
||||
ten_list: list[FakeTensor | Any],
|
||||
) -> list[FakeTensor | torch.Tensor | Any]:
|
||||
"""
|
||||
Attempt to convert fake tensors to a corresponding real tensor with the correct underlying storage by looking up
|
||||
the FakeTensorMode meta to real storage mapping. On failure to find the storage mapping, the FakeTensor will
|
||||
remain in the list.
|
||||
|
||||
Note: this is not currently optimized (makes copies of the meta converter internal dictionaries)
|
||||
"""
|
||||
|
||||
fake_tensor = next(
|
||||
(item for item in ten_list if isinstance(item, FakeTensor)), None
|
||||
)
|
||||
if fake_tensor is None:
|
||||
return ten_list
|
||||
|
||||
fake_mode = fake_tensor.fake_mode
|
||||
meta_converter = fake_mode.fake_tensor_converter.meta_converter
|
||||
desc = meta_converter.describer
|
||||
|
||||
storage_to_key = {v: k for k, v in meta_converter.storage_memo.items()}
|
||||
key_to_real_storage = {v: k for k, v in desc.lookup_storage.items()}
|
||||
out = []
|
||||
for t in ten_list:
|
||||
if not isinstance(t, FakeTensor) or t.layout != torch.strided:
|
||||
out.append(t)
|
||||
continue
|
||||
|
||||
key = storage_to_key.get(t.untyped_storage())
|
||||
real_storage = None if key is None else key_to_real_storage.get(key)
|
||||
if real_storage is None:
|
||||
out.append(t)
|
||||
continue
|
||||
|
||||
unhinted = False
|
||||
|
||||
def map_symint(s: torch.SymInt | int) -> int:
|
||||
nonlocal unhinted
|
||||
if not isinstance(s, torch.SymInt):
|
||||
return s
|
||||
unhinted = unhinted if not unhinted else s.node.has_hint()
|
||||
return s.node.hint
|
||||
|
||||
stor_offset = map_symint(t.storage_offset())
|
||||
size = [map_symint(s) for s in t.shape]
|
||||
stride = [map_symint(s) for s in t.stride()]
|
||||
|
||||
if unhinted:
|
||||
out.append(t)
|
||||
continue
|
||||
|
||||
new_tensor = torch.empty(
|
||||
[],
|
||||
dtype=t.dtype,
|
||||
device=t.device,
|
||||
)
|
||||
new_tensor.set_(
|
||||
real_storage,
|
||||
storage_offset=stor_offset,
|
||||
size=size,
|
||||
stride=stride,
|
||||
)
|
||||
out.append(new_tensor.clone())
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _check_fake_real_tensors(
|
||||
real_out: torch.Tensor,
|
||||
fake_out: FakeTensor,
|
||||
context: str = "",
|
||||
sizes: bool = True,
|
||||
strides: bool = False,
|
||||
storage_offset: bool = True,
|
||||
requires_grad: bool = True,
|
||||
) -> None:
|
||||
if requires_grad:
|
||||
if real_out.requires_grad != fake_out.requires_grad:
|
||||
raise MetadataMismatchError(
|
||||
f"{context} mismatched requires_grad-ness of outputs. "
|
||||
f"This usually means that you have added autograd support "
|
||||
f"for your operator at a dispatch key other than Autograd, "
|
||||
f"which will lead to problems"
|
||||
)
|
||||
|
||||
if torch._C._has_storage(real_out):
|
||||
r_offset = real_out.storage_offset()
|
||||
f_offset = fake_out.storage_offset()
|
||||
if r_offset != f_offset:
|
||||
raise MetadataMismatchError(f"{context} mismatched storage offset")
|
||||
|
||||
torch._prims.utils.compare_tensor_meta(
|
||||
real_out,
|
||||
fake_out,
|
||||
check_sizes=sizes,
|
||||
check_strides=strides,
|
||||
allow_rhs_unbacked=True,
|
||||
)
|
||||
|
||||
|
||||
class CrossRefFakeMode(TorchDispatchMode):
|
||||
def __init__(
|
||||
self,
|
||||
ignore_op_fn: Callable[[OpOverload], bool] | None = None,
|
||||
*,
|
||||
check_strides: bool = True,
|
||||
check_aliasing: bool = True,
|
||||
only_check_ops_with_meta: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.ignore_op_fn = (
|
||||
ignore_op_fn if ignore_op_fn is not None else lambda fn: False
|
||||
)
|
||||
self.check_strides = check_strides
|
||||
self.check_aliasing = check_aliasing
|
||||
self.only_check_ops_with_meta = only_check_ops_with_meta
|
||||
|
||||
def __torch_dispatch__(
|
||||
self,
|
||||
func: OpOverload,
|
||||
types: Sequence[type],
|
||||
args: Sequence[object] = (),
|
||||
kwargs: Mapping[str, object] | None = None,
|
||||
) -> object:
|
||||
kwargs = kwargs or {}
|
||||
|
||||
fake_r = None
|
||||
fake_args: Sequence[object] = ()
|
||||
fake_kwargs: Mapping[str, object] = {}
|
||||
|
||||
# empty_like excluded for now due to sparse complex
|
||||
# aten._to_dense.default this one is getting called with csc
|
||||
if (
|
||||
func
|
||||
not in (
|
||||
aten.lift_fresh.default,
|
||||
aten.lift_fresh_copy.default,
|
||||
aten.set_.source_Storage_storage_offset,
|
||||
)
|
||||
and not self.ignore_op_fn(func)
|
||||
and (
|
||||
not self.only_check_ops_with_meta
|
||||
or torch._subclasses.fake_impls.has_meta(func)
|
||||
)
|
||||
and torch.Tag.dynamic_output_shape not in func.tags
|
||||
and torch.Tag.inplace_view not in func.tags
|
||||
and torch.Tag.data_dependent_output not in func.tags
|
||||
):
|
||||
# Do not import symbolic_shapes at the top of the module as it imports sympy and that's slow
|
||||
from torch.fx.experimental.symbolic_shapes import ShapeEnv
|
||||
|
||||
try:
|
||||
# TODO: enable_python_dispatcher() here
|
||||
with FakeTensorMode(shape_env=ShapeEnv()) as fake_mode:
|
||||
fake_args, fake_kwargs = pytree.tree_map_only(
|
||||
torch.Tensor,
|
||||
functools.partial(fake_mode.from_tensor, static_shapes=True),
|
||||
(args, kwargs),
|
||||
)
|
||||
with warnings.catch_warnings():
|
||||
fake_r = func(*fake_args, **fake_kwargs)
|
||||
except UnsupportedFakeTensorException:
|
||||
pass
|
||||
|
||||
context = (
|
||||
f"When comparing the output of {func} on FakeTensor and concrete Tensors, "
|
||||
f"found"
|
||||
)
|
||||
r = func(*args, **kwargs)
|
||||
if fake_r is not None:
|
||||
r_flat = pytree.tree_leaves(r)
|
||||
f_flat = pytree.tree_leaves(fake_r)
|
||||
if len(f_flat) != len(r_flat):
|
||||
raise AssertionError(
|
||||
f"{context} mismatch in number of returns {len(f_flat)} != {len(r_flat)}"
|
||||
)
|
||||
|
||||
if self.check_aliasing:
|
||||
_check_alias_info(
|
||||
context, r, (args, kwargs), fake_r, (fake_args, fake_kwargs)
|
||||
)
|
||||
|
||||
for idx, (r_out, f_out) in enumerate(
|
||||
zip(pytree.tree_leaves(r), pytree.tree_leaves(fake_r))
|
||||
):
|
||||
r_is_ten = isinstance(r_out, torch.Tensor)
|
||||
if r_is_ten != isinstance(f_out, torch.Tensor):
|
||||
raise AssertionError(
|
||||
f"{context} mismatched number of tensor outputs"
|
||||
)
|
||||
if r_is_ten:
|
||||
try:
|
||||
_check_fake_real_tensors(
|
||||
r_out,
|
||||
f_out,
|
||||
sizes=True,
|
||||
strides=self.check_strides,
|
||||
storage_offset=True,
|
||||
requires_grad=True,
|
||||
)
|
||||
except Exception as e:
|
||||
if is_sdpa_error(func, idx, e):
|
||||
continue
|
||||
error_message = (
|
||||
f"{context} mismatched tensor metadata: {e}"
|
||||
if len(r_flat) == 1
|
||||
else f"{context} mismatched tensor metadata for output[{idx}]: {e}"
|
||||
)
|
||||
raise MetadataMismatchError(error_message) from e
|
||||
return r
|
||||
@@ -0,0 +1,894 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import warnings
|
||||
import weakref
|
||||
from abc import ABC, abstractmethod
|
||||
from contextlib import AbstractContextManager
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import builtins
|
||||
from collections.abc import Callable, Generator, Sequence
|
||||
from types import TracebackType
|
||||
|
||||
from torch._functorch.pyfunctorch import FunctionalizeInterpreter
|
||||
from torch._ops import OpOverload
|
||||
|
||||
import torch
|
||||
import torch.fx.traceback as fx_traceback
|
||||
import torch.utils._pytree as pytree
|
||||
from torch._C import _functionalization_reapply_views_tls as _reapply_views
|
||||
from torch._ops import _get_dispatch_mode_pre_dispatch, TorchBindOpOverload
|
||||
from torch._subclasses.meta_utils import is_sparse_any
|
||||
from torch.utils._python_dispatch import (
|
||||
_detect_infra_mode,
|
||||
_disable_infra_mode,
|
||||
autograd_would_have_decomposed,
|
||||
return_and_correct_aliasing,
|
||||
TorchDispatchMode,
|
||||
)
|
||||
|
||||
|
||||
not_implemented_log = torch._logging.getArtifactLogger(__name__, "not_implemented")
|
||||
|
||||
|
||||
def _has_unrecognized_tensor_types(types: Sequence[type]) -> bool:
|
||||
unrecognized_types = [
|
||||
t
|
||||
for t in types
|
||||
if t not in (torch.Tensor, torch._subclasses.FakeTensor, FunctionalTensor)
|
||||
]
|
||||
if unrecognized_types:
|
||||
not_implemented_log.debug(
|
||||
"FunctionalTensor unrecognized subclass(es): %s", unrecognized_types
|
||||
)
|
||||
return bool(unrecognized_types)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=512)
|
||||
def _can_decompose_fast(
|
||||
func: OpOverload, export: bool, pre_dispatch: bool
|
||||
) -> bool | None:
|
||||
"""Fast path for _can_decompose that depends only on (func, export, pre_dispatch).
|
||||
|
||||
Returns True/False for a definitive answer, or None to fall through
|
||||
to the slow path (autograd_would_have_decomposed).
|
||||
"""
|
||||
if export and func is torch.ops.aten.dropout.default:
|
||||
return False
|
||||
|
||||
from torch._decomp import _should_decompose_because_unsafe_op
|
||||
|
||||
if _should_decompose_because_unsafe_op(func):
|
||||
return True
|
||||
|
||||
alias_info_present = any(arg.alias_info for arg in func._schema.arguments)
|
||||
if alias_info_present or func._schema.is_mutable:
|
||||
return True
|
||||
|
||||
if export:
|
||||
if pre_dispatch:
|
||||
if func.namespace not in ("aten", "prim") and func._can_decompose():
|
||||
warnings.warn(
|
||||
f"At pre-dispatch tracing, we assume that any custom op marked with "
|
||||
f"CompositeImplicitAutograd and have functional schema are safe to not decompose. "
|
||||
f"Found {func} to be one such op.",
|
||||
stacklevel=3,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _assert_functionalize_not_active(msg: str) -> None:
|
||||
is_included = torch._C._dispatch_tls_is_dispatch_key_included(
|
||||
torch._C.DispatchKey.Functionalize
|
||||
)
|
||||
is_excluded = torch._C._dispatch_tls_is_dispatch_key_excluded(
|
||||
torch._C.DispatchKey.Functionalize
|
||||
)
|
||||
if not is_excluded and is_included:
|
||||
raise AssertionError(msg)
|
||||
|
||||
|
||||
# NOTE Some special handling for tensor conversion during export is needed.
|
||||
# Normally, when tracing through the model with tensor.to(), the maybe-aliasing
|
||||
# relationship between input and output tensors will be baked into the graph.
|
||||
# For example, if we got a tensor with device cpu and call tensor.to("cpu"),
|
||||
# it will become a no-op in the graph. For a whole graph capture, this is not
|
||||
# sound so we need to do something different. Instead, in export we will try to
|
||||
# preserve the tensor conversion by forcing a non-semantic-breaking aten::_to_copy
|
||||
# operator to be traced in the graph, and subsequently banning mutations on all
|
||||
# such converted tensors.
|
||||
# In addition to patching .to() method call in functionalization, we will have to
|
||||
# patch other similar methods like float() and cpu(), because they intentionally
|
||||
# don't fall back to .to() methods, but have the same behavior as .to() according to
|
||||
# pytorch document. https://pytorch.org/docs/stable/generated/torch.Tensor.float.html
|
||||
# thus we simply force them to go through .to() call.
|
||||
def _conversion_method_template(**extra_kwargs: Any) -> Callable[..., Any]:
|
||||
def _(self: FunctionalTensor, *args: Any, **kwargs: Any) -> Any:
|
||||
return self.to(*args, **{**kwargs, **extra_kwargs})
|
||||
|
||||
return _
|
||||
|
||||
|
||||
class FunctionalTensor(torch.Tensor):
|
||||
"""
|
||||
Functional tensors represent tensors that will remove mutations
|
||||
from a program. If you perform a mutable operation on a functional tensor,
|
||||
it will re-dispatch to the functional variant of that operation.
|
||||
|
||||
Historically, functionalization is implemented in C++ in the dispatcher.
|
||||
This class is a lightweight python shim around the C++ functionalization logic.
|
||||
|
||||
FunctionalTensor is required to be used with a corresponding
|
||||
FunctionalTensorMode active, because it relies
|
||||
on using the mode for dispatch (which can properly handle factory functions).
|
||||
"""
|
||||
|
||||
elem: torch.Tensor
|
||||
# Indicates to our torch_dispatch dispatching infra that
|
||||
# this is an "infra" mode with lower dispatching precedence.
|
||||
_mode_key = torch._C._TorchDispatchModeKey.FUNCTIONAL
|
||||
|
||||
# Note: The reason we add these extra keys to our FunctionalTensor subclass
|
||||
# is to mirror the behavior of C++ functionalization (we can choose to change this
|
||||
# later, as long as it doesn't break anything).
|
||||
# FunctionalTensorWrapper copies **all** dispatch keys from the inner tensor
|
||||
# to the wrapper, excluding functorch and python dispatch keys.
|
||||
# Here I'm trying to reuse the keyset the functorch wrapper subclasses copy,
|
||||
# except that they don't include ZeroTensor so I'm manually adding it in.
|
||||
_extra_dispatch_keys = torch._C._additional_keys_to_prop_for_wrapper_tensors.add(
|
||||
torch._C.DispatchKey.ZeroTensor
|
||||
)
|
||||
|
||||
# These are all aten ops that correspond to metadata queries.
|
||||
# We want FunctionalTensor to be able to handle them directly.
|
||||
metadata_fns = frozenset(
|
||||
{
|
||||
torch.ops.aten.is_contiguous.default,
|
||||
torch.ops.aten.is_contiguous.memory_format,
|
||||
torch.ops.aten.is_strides_like_format.default,
|
||||
torch.ops.aten.is_non_overlapping_and_dense.default,
|
||||
torch.ops.aten.size.default,
|
||||
torch.ops.aten.sym_size.default,
|
||||
torch.ops.aten.stride.default,
|
||||
torch.ops.aten.sym_stride.default,
|
||||
torch.ops.aten.storage_offset.default,
|
||||
torch.ops.aten.sym_storage_offset.default,
|
||||
torch.ops.aten.numel.default,
|
||||
torch.ops.aten.sym_numel.default,
|
||||
torch.ops.aten.dim.default,
|
||||
torch.ops.prim.device.default,
|
||||
}
|
||||
)
|
||||
|
||||
# Used by auto_functionalize to determine base of tensors during inference mode.
|
||||
_inference_mode_base: FunctionalTensor | None = None
|
||||
|
||||
def __new__(cls, elem: torch.Tensor, mode: FunctionalTensorMode) -> Self:
|
||||
if not torch._is_functional_tensor(elem):
|
||||
raise AssertionError("elem must be a functional tensor")
|
||||
|
||||
# In general, we'd like our functional tensor subclass to only be in charge of functionalization,
|
||||
# and defer to the inner subclass for all other functionality.
|
||||
# Example: If our inner tensor is a ZeroTensor, we would want to defer running the ZeroTensor fallback
|
||||
# until after we redispatch to our inner ZeroTensor.
|
||||
# However, there are a few keys that we need to mirror between the inner and outer tensors.
|
||||
# Conjugate
|
||||
# Negative
|
||||
# Why? These keys are used to test metadata queries, like `.is_conj()` and `.is_neg()`.
|
||||
# We **need** calls to is_conj() to return the same thing on the outer and inner tensors,
|
||||
# Because user code / framework code that branches like so needs to do the same thing
|
||||
# when it sees the outer FunctionalTensor:
|
||||
# if (x.is_conj()) {
|
||||
# return at::view_as_real(x.resolve_conj());
|
||||
# } else {
|
||||
# return at::view_as_real(x);
|
||||
# }
|
||||
extra_dispatch_keys = (
|
||||
FunctionalTensor._extra_dispatch_keys & torch._C._dispatch_keys(elem)
|
||||
)
|
||||
|
||||
out = torch.Tensor._make_wrapper_subclass(
|
||||
# TODO: right now, _make_wrapper_subclass's dynamic shape interaction is not great.
|
||||
# Calling the overload that has kwargs causes us to go down the first overload path,
|
||||
# which will **always** specialize sizes.
|
||||
# We should probably eventually fix this so that the first overload can just handle dynamic shapes.
|
||||
cls,
|
||||
elem.shape, # sizes
|
||||
elem.stride() if not is_sparse_any(elem) else None, # strides
|
||||
(
|
||||
elem.storage_offset() if not is_sparse_any(elem) else None
|
||||
), # storage_offset
|
||||
None, # memory_format
|
||||
elem.dtype, # dtype
|
||||
elem.layout, # layout
|
||||
elem.device, # device
|
||||
False, # pin_memory
|
||||
elem.requires_grad, # requires_grad
|
||||
None, # dispatch_sizes_strides_policy
|
||||
False, # dispatch_device
|
||||
False, # dispatch_layout
|
||||
extra_dispatch_keys, # _extra_dispatch_keys
|
||||
)
|
||||
torch._C._set_throw_on_mutable_data_ptr(out)
|
||||
out.elem = elem
|
||||
|
||||
if (
|
||||
torch._export.config.enable_auto_functionalized_v2_for_export
|
||||
and torch.is_inference_mode_enabled()
|
||||
and torch._inductor.config.enable_auto_functionalized_v2
|
||||
):
|
||||
if out.is_base_tensor():
|
||||
out._inference_mode_base = None
|
||||
# This assumes that the FunctionalTensor.elem does not change its storage after this point.
|
||||
# Otherwise this would be invalid.
|
||||
mode._storage_to_base[out.elem.untyped_storage()] = out
|
||||
else:
|
||||
out._inference_mode_base = mode._storage_to_base[
|
||||
out.elem.untyped_storage()
|
||||
]
|
||||
if out._inference_mode_base is None:
|
||||
raise AssertionError("out._inference_mode_base must not be None")
|
||||
return out
|
||||
|
||||
def __torch_dispatch__( # type: ignore[override]
|
||||
self,
|
||||
func: OpOverload,
|
||||
types: Sequence[type],
|
||||
args: tuple[Any, ...] = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
if _has_unrecognized_tensor_types(types):
|
||||
return NotImplemented
|
||||
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
# FunctionalTensor needs to plumb all metadata requests to the inner tensor.
|
||||
# In theory we don't have to do this - but if we want to service metadata requests here,
|
||||
# we need to carefully make sure all metadata is accurate (including metadata mutations)
|
||||
if func in FunctionalTensor.metadata_fns:
|
||||
# All metadata accesses should be plumbed to the inner tensor, that way we don't have to worry
|
||||
# about the problem of keeping metadata in sync between the wrapper and inner tensor.
|
||||
# This also alleviates us from having to manually handle metadata mutations on the wrapper.
|
||||
if len(kwargs) != 0:
|
||||
raise AssertionError("kwargs must be empty for metadata functions")
|
||||
if func in [
|
||||
torch.ops.aten.is_strides_like_format.default,
|
||||
torch.ops.aten.is_contiguous.memory_format,
|
||||
]:
|
||||
if len(args) != 2 or not isinstance(args[0], FunctionalTensor):
|
||||
raise AssertionError("Expected 2 args with FunctionalTensor first")
|
||||
return func(torch._from_functional_tensor(args[0].elem), args[1])
|
||||
if len(args) != 1 or not isinstance(args[0], FunctionalTensor):
|
||||
raise AssertionError("Expected 1 arg with FunctionalTensor")
|
||||
|
||||
return func(torch._from_functional_tensor(args[0].elem))
|
||||
# Originally I tried to implement my subclass without giving it a torch_dispatch, but I gave up:
|
||||
# - _make_wrapper_subclass requires a __torch_dispatch__
|
||||
# - If we want to use _make_subclass(), we have a problem: the subclass will share a TensorImpl with the inner tensor,
|
||||
# which is of type FunctionalTensorWrapper! We explicitly do not want our wrapper to be a FunctionalTensorWrapper.
|
||||
# - If we use the default tensor.__new__(), we have another problem: it returns inner_tensor.alias(),
|
||||
# which causes every subclass created above autograd to have autograd view metadata
|
||||
# (in addition to also being a FunctionalTensorWrapper).
|
||||
raise RuntimeError(
|
||||
"Attempting to use FunctionalTensor on its own. Instead, please use it with a corresponding FunctionalTensorMode()"
|
||||
)
|
||||
|
||||
def __repr__(self, *, tensor_contents: object | None = None) -> str:
|
||||
return f"FunctionalTensor({repr(self.elem)})"
|
||||
|
||||
@staticmethod
|
||||
def to_functional(x: torch.Tensor) -> FunctionalTensor:
|
||||
# We will do the wrapping for the user.
|
||||
|
||||
if torch._is_functional_tensor(x):
|
||||
raise AssertionError("x must not already be a functional tensor")
|
||||
# The only autograd metadata we care about on the FunctionalTensor is:
|
||||
# - requires_grad (so autograd runs)
|
||||
# - is_leaf (so that mutations on graph inputs that are not leaves are allowed by the autograd engine)
|
||||
# this is handled by FunctionalTensor.to_functional
|
||||
x_functional = torch._to_functional_tensor(x)
|
||||
# Technically the FunctionalTensormode here is unnecessary,
|
||||
# but it avoids spurious NotImplemented logs during `ProxyTorchDispatchMode` tracing.
|
||||
# _mirror_autograd_meta_to queries tensor sizes,
|
||||
# and otherwise the sym_size() call will go to the proxy mode before hitting
|
||||
# FunctionalTensor.__torch_dispatch__
|
||||
|
||||
functional_mode = _detect_infra_mode(torch._C._TorchDispatchModeKey.FUNCTIONAL)
|
||||
if functional_mode is None:
|
||||
raise AssertionError("functional_mode must not be None")
|
||||
|
||||
with functional_mode:
|
||||
torch._mirror_autograd_meta_to(x, x_functional) # type: ignore[attr-defined]
|
||||
out = FunctionalTensor(x_functional, functional_mode)
|
||||
torch._mirror_autograd_meta_to(x_functional, out) # type: ignore[attr-defined]
|
||||
return out
|
||||
|
||||
def from_functional(self) -> torch.Tensor:
|
||||
torch._sync(self)
|
||||
return torch._from_functional_tensor(self.elem)
|
||||
|
||||
def is_base_tensor(self) -> bool:
|
||||
return torch._is_functional_tensor_base(self.elem)
|
||||
|
||||
def replace_(self, output: torch.Tensor) -> None:
|
||||
torch._functionalize_replace(self.elem, output)
|
||||
|
||||
def commit_update(self) -> None:
|
||||
torch._functionalize_commit_update(self.elem)
|
||||
|
||||
def sync(self) -> None:
|
||||
torch._functionalize_sync(self.elem)
|
||||
|
||||
def mark_mutation_hidden_from_autograd(self) -> None:
|
||||
torch._functionalize_mark_mutation_hidden_from_autograd(self.elem)
|
||||
|
||||
def tolist(self) -> Any:
|
||||
if self.elem.dim() == 0:
|
||||
return self.elem.item()
|
||||
elif self.elem.dim() == 1:
|
||||
return [elem.item() for elem in self.elem]
|
||||
else:
|
||||
return [elem.tolist() for elem in self.elem]
|
||||
|
||||
def to(self, *args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
if _detect_infra_mode(torch._C._TorchDispatchModeKey.FUNCTIONAL).export:
|
||||
torch.ops.aten._assert_tensor_metadata(
|
||||
self,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
layout=self.layout,
|
||||
)
|
||||
return super().to(*args, **kwargs)
|
||||
|
||||
# pyrefly: ignore[bad-override]
|
||||
def cuda(
|
||||
self, device: torch.device | int | str | None = None, *args: Any, **kwargs: Any
|
||||
) -> torch.Tensor:
|
||||
device = device or torch.cuda.current_device()
|
||||
if len(args) > 0:
|
||||
return self.to(device, *args, **kwargs)
|
||||
else:
|
||||
return self.to(device=device, **kwargs)
|
||||
|
||||
char = _conversion_method_template(dtype=torch.int8)
|
||||
cpu = _conversion_method_template(device=torch.device("cpu"))
|
||||
bfloat16 = _conversion_method_template(dtype=torch.bfloat16)
|
||||
byte = _conversion_method_template(dtype=torch.uint8)
|
||||
double = _conversion_method_template(dtype=torch.float64)
|
||||
float = _conversion_method_template(dtype=torch.float32)
|
||||
bool = _conversion_method_template(dtype=torch.bool)
|
||||
half = _conversion_method_template(dtype=torch.float16)
|
||||
int = _conversion_method_template(dtype=torch.int32)
|
||||
long = _conversion_method_template(dtype=torch.int64)
|
||||
|
||||
# TODO(sparse-team): fixes #133174 but can we do without the relay?
|
||||
def to_dense(
|
||||
self,
|
||||
dtype: torch.dtype | None = None,
|
||||
*,
|
||||
masked_grad: builtins.bool | None = None,
|
||||
) -> torch.Tensor:
|
||||
return self.elem.to_dense()
|
||||
|
||||
@property
|
||||
# pyrefly: ignore[bad-override]
|
||||
def layout(self) -> torch.layout:
|
||||
return self.elem.layout
|
||||
|
||||
def __bool__(self) -> builtins.bool:
|
||||
return bool(self.item())
|
||||
|
||||
|
||||
class FunctionalTensorMode(TorchDispatchMode):
|
||||
def __init__(
|
||||
self,
|
||||
pre_dispatch: bool = False,
|
||||
export: bool = False,
|
||||
_allow_token_discovery: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.export = export
|
||||
self.is_on_stack = False
|
||||
self.enter_stack = []
|
||||
# Indicates to our torch_dispatch dispatching infra that
|
||||
# this is an "infra" mode with lower dispatching precedence.
|
||||
self._mode_key = torch._C._TorchDispatchModeKey.FUNCTIONAL
|
||||
self.pre_dispatch = pre_dispatch
|
||||
# This will be turned off later for pre-dispatch functionalization
|
||||
self._dispatch_key = torch._C.DispatchKey.PreDispatch if pre_dispatch else None # type: ignore[attr-defined]
|
||||
# Map of effect type (ex. _EffectType.ORDERED) to a token. The tokens help keep
|
||||
# track of the ordering between side effectful operations.
|
||||
self._tokens: dict[Any, torch.Tensor] = {}
|
||||
|
||||
# Filled after forward tracing.
|
||||
self._tokens_forward_output: dict[Any, torch.Tensor] = {}
|
||||
|
||||
# Functionalization runs twice in AOTAutograd, once in
|
||||
# `run_functionalized_fw_and_collect_metadata` to collect metadata to
|
||||
# see which tensors need to be functionalized and discover how many
|
||||
# tokens we need, and another time in `make_fx` which does the actual
|
||||
# tracing to replace ops with their functional variants and handling
|
||||
# side-effectful ops. In the second stage there should be no token
|
||||
# discovery. This flag distinguishes between the two stages.
|
||||
self._allow_token_discovery = _allow_token_discovery
|
||||
|
||||
self._storage_to_base: weakref.WeakKeyDictionary[
|
||||
torch.storage.UntypedStorage, FunctionalTensor | None
|
||||
] = weakref.WeakKeyDictionary()
|
||||
|
||||
# No-op if FunctionalTensorMode is already in use
|
||||
def __enter__(self) -> Self:
|
||||
def _get_prev_mode() -> FunctionalTensorMode | None:
|
||||
if self._dispatch_key == torch._C.DispatchKey.PreDispatch:
|
||||
return _get_dispatch_mode_pre_dispatch(
|
||||
torch._C._TorchDispatchModeKey.FUNCTIONAL
|
||||
)
|
||||
return torch._C._get_dispatch_mode(
|
||||
torch._C._TorchDispatchModeKey.FUNCTIONAL
|
||||
)
|
||||
|
||||
if _get_prev_mode() is None:
|
||||
self.enter_stack.append(True)
|
||||
return super().__enter__()
|
||||
else:
|
||||
self.enter_stack.append(False)
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: TracebackType | None,
|
||||
) -> None:
|
||||
is_on_stack = self.enter_stack.pop()
|
||||
if is_on_stack:
|
||||
super().__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
def __torch_dispatch__(
|
||||
self,
|
||||
func: OpOverload,
|
||||
types: Sequence[type],
|
||||
args: tuple[Any, ...] = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
if _has_unrecognized_tensor_types(types):
|
||||
return NotImplemented
|
||||
|
||||
if (
|
||||
func not in FunctionalTensor.metadata_fns
|
||||
and self._can_decompose(func, args, kwargs)
|
||||
# Not all funcs from __torch_dispatch__ are actual dispatcher ops,
|
||||
# e.g. prim.device
|
||||
and torch._C._dispatch_has_kernel(func.name())
|
||||
):
|
||||
with self:
|
||||
r = func.decompose(*args, **kwargs)
|
||||
if r is not NotImplemented:
|
||||
return r
|
||||
|
||||
def wrap(x: object) -> object:
|
||||
# Only wrap our outputs in subclasses if the inner functionalization call
|
||||
# also wrapped outputs into FunctionalTensorWrappers.
|
||||
# When can this happen? e.g. `torch.div(2, 2)`
|
||||
if isinstance(x, FunctionalTensor):
|
||||
raise AssertionError("x must not be a FunctionalTensor in wrap()")
|
||||
if isinstance(x, torch.Tensor) and torch._is_functional_tensor(x):
|
||||
return FunctionalTensor(x, self)
|
||||
return x
|
||||
|
||||
def unwrap(x: FunctionalTensor) -> torch.Tensor:
|
||||
return x.elem
|
||||
|
||||
from torch._higher_order_ops.auto_functionalize import (
|
||||
can_auto_functionalize,
|
||||
do_auto_functionalize,
|
||||
do_auto_functionalize_v2,
|
||||
)
|
||||
|
||||
if can_auto_functionalize(
|
||||
func
|
||||
) and not torch._C._dispatch_has_kernel_for_dispatch_key(
|
||||
func.name(), torch._C.DispatchKey.Functionalize
|
||||
):
|
||||
import torch._export.config as export_config
|
||||
import torch._inductor.config as inductor_config
|
||||
|
||||
if torch.compiler.is_exporting():
|
||||
if export_config.enable_auto_functionalized_v2_for_export:
|
||||
return do_auto_functionalize_v2(self, func, args, kwargs)
|
||||
|
||||
return do_auto_functionalize(self, func, args, kwargs)
|
||||
|
||||
if inductor_config.enable_auto_functionalized_v2:
|
||||
return do_auto_functionalize_v2(self, func, args, kwargs)
|
||||
return do_auto_functionalize(self, func, args, kwargs)
|
||||
|
||||
from torch._higher_order_ops.effects import handle_effects, has_effects
|
||||
|
||||
if has_effects(func):
|
||||
if torch._C._dispatch_has_kernel_for_dispatch_key(
|
||||
func.name(), torch._C.DispatchKey.Functionalize
|
||||
):
|
||||
raise AssertionError(
|
||||
f"func {func.name()} with effects should not have a kernel for Functionalize dispatch key"
|
||||
)
|
||||
return handle_effects(
|
||||
self._allow_token_discovery, self._tokens, func, args, kwargs
|
||||
)
|
||||
|
||||
args_unwrapped, kwargs_unwrapped = pytree.tree_map_only(
|
||||
FunctionalTensor, unwrap, (args, kwargs)
|
||||
)
|
||||
|
||||
# Expectation: functionalization should not **already** be enabled above our mode.
|
||||
# Why would that be bad? when we return a FunctionalTensor here, we don't want functionalization
|
||||
# to run above this mode and further wrap that output in **another** C++ FunctionalTensorWrapper.
|
||||
_assert_functionalize_not_active(
|
||||
"Functionalization should not already be enabled above this mode"
|
||||
)
|
||||
include_to_set = (
|
||||
torch._C._dispatch_tls_local_include_set()
|
||||
| torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
|
||||
)
|
||||
exclude_to_set = (
|
||||
torch._C._dispatch_tls_local_exclude_set().remove(
|
||||
torch._C.DispatchKey.Functionalize
|
||||
)
|
||||
- FunctionalTensor._extra_dispatch_keys
|
||||
)
|
||||
|
||||
if isinstance(func, TorchBindOpOverload):
|
||||
# When the function is a TorchBindOpOverload, meaning some of the
|
||||
# inputs are FakeScriptObjects, we need to skip c++ dispatcher and
|
||||
# dispatch in python because C++ dispatcher will check the schema
|
||||
# and cannot recognize FakeScriptObject.
|
||||
ctx = PythonFunctionalizeAPI()
|
||||
fully_unwrapped_args = ctx.unwrap_tensors(args)
|
||||
fully_unwrapped_kwargs = ctx.unwrap_tensors(
|
||||
kwargs # pyrefly: ignore[bad-argument-type]
|
||||
)
|
||||
outs_unwrapped = func(
|
||||
*fully_unwrapped_args,
|
||||
**fully_unwrapped_kwargs,
|
||||
)
|
||||
outs_wrapped = ctx.wrap_tensors(outs_unwrapped)
|
||||
else:
|
||||
# All we want to do here is reuse the existing C++ functionalization logic.
|
||||
# This requires swizzling our TLS dispatch keys so that the Functionalize key is active.
|
||||
with torch._C._ForceDispatchKeyGuard(include_to_set, exclude_to_set):
|
||||
try:
|
||||
# By default for python functionalization (for AOTAutograd), we reapply views.
|
||||
old_apply_views = torch._functionalize_enable_reapply_views(True) # type: ignore[attr-defined]
|
||||
|
||||
# Sometimes these functions cannot be directly dispatched to functionalize key
|
||||
# because args are sometimes not functional tensors for some reason?
|
||||
if func in FunctionalTensor.metadata_fns:
|
||||
outs_unwrapped = func(*args_unwrapped, **kwargs_unwrapped)
|
||||
outs_wrapped = pytree.tree_map_only(
|
||||
torch.Tensor, wrap, outs_unwrapped
|
||||
)
|
||||
else:
|
||||
self._sync_view_replay_annotations(args, kwargs)
|
||||
|
||||
# When we dispatch to the C++ functionalization kernel, we might need to jump back to the
|
||||
# PreDispatch mode stack afterwards, to handle any other PreDispatch modes underneath
|
||||
# FunctionalTensorMode. If we call func() directly, we would need to exclude PreDispatch
|
||||
# from the TLS in order to avoid infinite looping, but this would prevent us from coming
|
||||
# back to PreDispatch later
|
||||
outs_unwrapped = func._op_dk(
|
||||
torch._C.DispatchKey.Functionalize,
|
||||
*args_unwrapped,
|
||||
**kwargs_unwrapped,
|
||||
)
|
||||
|
||||
if self.export:
|
||||
if func is torch.ops.aten.dropout.default:
|
||||
torch._freeze_functional_tensor(outs_unwrapped) # type: ignore[attr-defined]
|
||||
outs_wrapped = pytree.tree_map_only(
|
||||
torch.Tensor, wrap, outs_unwrapped
|
||||
)
|
||||
finally:
|
||||
torch._disable_functionalization()
|
||||
torch._functionalize_enable_reapply_views(old_apply_views) # type: ignore[attr-defined]
|
||||
|
||||
_assert_functionalize_not_active(
|
||||
"Functionalization should not already be enabled above this mode after dispatch"
|
||||
)
|
||||
|
||||
if (
|
||||
# If no outputs are our functional subclass, then don't try to fix up aliasing
|
||||
not any(
|
||||
isinstance(x, FunctionalTensor)
|
||||
for x in pytree.tree_leaves(outs_wrapped)
|
||||
)
|
||||
# Since lift_fresh lifts its argument into a functional tensor, we can skip the
|
||||
# aliasing correction step. Otherwise, we would be setting the storage of a
|
||||
# lifted tensor to that of an unlifted tensor.
|
||||
# Ref: https://github.com/pytorch/pytorch/issues/111506
|
||||
or func is torch.ops.aten.lift_fresh.default
|
||||
):
|
||||
return outs_wrapped
|
||||
# for metadata mutations, need to manually mutate the metadata of the FunctionalTensor wrapper
|
||||
if (
|
||||
torch.Tag.inplace_view in func.tags
|
||||
and func is not torch.ops.aten.set_.source_Tensor
|
||||
):
|
||||
with torch.utils._mode_utils.no_dispatch():
|
||||
func(*args, **kwargs)
|
||||
# Wrapper tensor subclasses do not have correct aliasing info! Use this util to manually correct the output aliasing.
|
||||
# inplace ops like `aten.add_()` are expected to return inputs **directly**, instead of creating fresh tensor objects.
|
||||
# Use this util to figure out the right thing to return.
|
||||
# If none of our inputs were wrapped, then we have no FunctionalTensor outputs that we need to fix up storages for.
|
||||
return return_and_correct_aliasing(func, args, kwargs, outs_wrapped)
|
||||
|
||||
def _sync_view_replay_annotations(
|
||||
self,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
"""Sync FunctionalTensor args so view replay uses correct fx node metadata.
|
||||
|
||||
When functionalization encounters a mutation, it handles aliases by lazily
|
||||
regenerating them at the first time they are next used. This is a problem when
|
||||
plumbing user annotations during tracing: we want view ops from view replay to
|
||||
have the same annotation the user specified on the original views. But view
|
||||
replay happens the next time the alias is used (e.g.
|
||||
second_op(alias_with_pending_mutation)), so the regenerated views would get the
|
||||
metadata for second_op instead.
|
||||
|
||||
To fix this, we remember the node metadata from the original views and globally
|
||||
set it when we lazily perform view replay. The globally set metadata will be
|
||||
used to populate the fx node created for the replayed operation.
|
||||
"""
|
||||
m = torch._C._get_dispatch_mode(torch._C._TorchDispatchModeKey.PROXY)
|
||||
if m is not None:
|
||||
for a in pytree.tree_leaves([args, kwargs]):
|
||||
if not isinstance(a, FunctionalTensor):
|
||||
continue
|
||||
unwrapped = torch._from_functional_tensor(a.elem)
|
||||
try:
|
||||
tracker_entry = m.tracer.tensor_tracker[unwrapped]
|
||||
except KeyError:
|
||||
raise RuntimeError(
|
||||
f"cannot find {unwrapped} in tensor_tracker"
|
||||
) from None
|
||||
curr_node = tracker_entry.proxy.node
|
||||
with fx_traceback.set_current_replay_node(curr_node):
|
||||
torch._sync(a)
|
||||
|
||||
def _can_decompose(
|
||||
self,
|
||||
func: OpOverload,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any],
|
||||
) -> bool:
|
||||
result = _can_decompose_fast(func, self.export, self.pre_dispatch)
|
||||
if result is not None:
|
||||
return result
|
||||
|
||||
# in normal torch.compile IR, we only decompose an op if autograd
|
||||
# would have decomposed it (NB: autograd may have been skipped if
|
||||
# we are in inference mode)
|
||||
# TODO: the flatten here can potentially be deduped with the
|
||||
# unwrapping pytree_map later
|
||||
flat_args_kwargs, _ = pytree.tree_flatten((args, kwargs))
|
||||
return autograd_would_have_decomposed(func, flat_args_kwargs)
|
||||
|
||||
@classmethod
|
||||
def is_infra_mode(cls) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def disable_functional_mode() -> Generator[None, None, None]:
|
||||
return _disable_infra_mode(torch._C._TorchDispatchModeKey.FUNCTIONAL)
|
||||
|
||||
|
||||
# This is similar to torch.func.functionalize, but:
|
||||
# - It uses FunctionalTensorMode, and FunctionalTensor (a python subclass).
|
||||
# One important advantage to using this mode is that it will let us
|
||||
# run functionalization underneath __torch_dispatch__,
|
||||
# which we need in AOTAutograd.
|
||||
# - Doing so means that it does not automatically compose with other
|
||||
# functorch transforms, since these transforms always run above __torch_dispatch__.
|
||||
# That's why this util lives here, and not in functorch.
|
||||
def dispatch_functionalize(
|
||||
func: Callable[..., Any], mode: FunctionalTensorMode = FunctionalTensorMode()
|
||||
) -> Callable[..., Any]:
|
||||
# TODO: pull these from aot autograd
|
||||
def to_fun(t: object) -> object:
|
||||
if isinstance(t, torch.Tensor):
|
||||
return FunctionalTensor.to_functional(t)
|
||||
return t
|
||||
|
||||
def from_fun(t: object) -> object:
|
||||
if not isinstance(t, FunctionalTensor):
|
||||
# quick sanity check
|
||||
if isinstance(t, torch.Tensor):
|
||||
if torch._is_functional_tensor(t):
|
||||
raise AssertionError(
|
||||
"Non-FunctionalTensor torch.Tensor should not be a functional tensor"
|
||||
)
|
||||
return t
|
||||
torch._sync(t)
|
||||
return torch._from_functional_tensor(t.elem)
|
||||
|
||||
def inner(*args: Any, **kwargs: Any) -> Any:
|
||||
disable_above = torch._C._ExcludeDispatchKeyGuard(
|
||||
torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
|
||||
)
|
||||
with disable_above, mode:
|
||||
func_args = pytree.tree_map_only(torch.Tensor, to_fun, args)
|
||||
func_kwargs = pytree.tree_map_only(torch.Tensor, to_fun, kwargs)
|
||||
func_outputs = func(*func_args, **func_kwargs)
|
||||
outputs = pytree.tree_map_only(FunctionalTensor, from_fun, func_outputs)
|
||||
|
||||
return outputs
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
class BaseFunctionalizeAPI(ABC):
|
||||
@abstractmethod
|
||||
def wrap_tensors(self, args: tuple[Any, ...]) -> tuple[Any, ...]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def unwrap_tensors(self, args: torch.Tensor | tuple[torch.Tensor, ...]) -> Any:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def functionalize(self, inner_f: Callable[..., Any]) -> Callable[..., Any]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def redispatch_to_next(self) -> AbstractContextManager[None]:
|
||||
pass
|
||||
|
||||
def replace(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor) -> None:
|
||||
torch._functionalize_replace(input_tensor, output_tensor)
|
||||
|
||||
def commit_update(self, tensor: torch.Tensor) -> None:
|
||||
torch._functionalize_commit_update(tensor)
|
||||
|
||||
def sync(self, tensor: torch.Tensor) -> None:
|
||||
torch._functionalize_sync(tensor)
|
||||
|
||||
def mark_mutation_hidden_from_autograd(self, tensor: torch.Tensor) -> None:
|
||||
torch._functionalize_mark_mutation_hidden_from_autograd(tensor)
|
||||
|
||||
|
||||
class PythonFunctionalizeAPI(BaseFunctionalizeAPI):
|
||||
def __init__(
|
||||
self, mode: FunctionalTensorMode | None = None, pre_dispatch: bool = False
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.mode = mode if mode else FunctionalTensorMode()
|
||||
self.pre_dispatch = pre_dispatch
|
||||
|
||||
def wrap_tensors(self, args: tuple[Any]) -> tuple[Any]:
|
||||
with self.mode:
|
||||
return torch.utils._pytree.tree_map_only(
|
||||
torch.Tensor, FunctionalTensor.to_functional, args
|
||||
)
|
||||
|
||||
def unwrap_tensors(
|
||||
self, args: torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor]
|
||||
) -> Any:
|
||||
return torch.utils._pytree.tree_map_only(
|
||||
FunctionalTensor, FunctionalTensor.from_functional, args
|
||||
)
|
||||
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def functionalize(self, inner_f: Callable) -> Callable:
|
||||
return dispatch_functionalize(inner_f, self.mode)
|
||||
|
||||
def redispatch_to_next(self) -> AbstractContextManager[None]:
|
||||
# [NOTE] We don't do anything here because at the time
|
||||
# we exercise this path, we would have already popped the
|
||||
# FunctionalTensorMode from mode stack. Since FunctionalTensorMode
|
||||
# is now stateful, it is better to explicitly pass in correct mode
|
||||
# directly instead of globally setting it.
|
||||
return contextlib.nullcontext()
|
||||
|
||||
@staticmethod
|
||||
def _check_cast_functional(tensor: torch.Tensor, name: str) -> FunctionalTensor:
|
||||
if not isinstance(tensor, FunctionalTensor):
|
||||
raise AssertionError(
|
||||
f"{name} must be a FunctionalTensor, got {type(tensor)}"
|
||||
)
|
||||
return tensor
|
||||
|
||||
def replace(self, input_tensor: torch.Tensor, output_tensor: torch.Tensor) -> None:
|
||||
ft = self._check_cast_functional(input_tensor, "input_tensor")
|
||||
if isinstance(output_tensor, FunctionalTensor):
|
||||
raise AssertionError("output_tensor must not be a FunctionalTensor")
|
||||
ft.replace_(output_tensor)
|
||||
|
||||
def commit_update(self, tensor: torch.Tensor) -> None:
|
||||
self._check_cast_functional(tensor, "tensor").commit_update()
|
||||
|
||||
def sync(self, tensor: torch.Tensor) -> None:
|
||||
self._check_cast_functional(tensor, "tensor").sync()
|
||||
|
||||
def mark_mutation_hidden_from_autograd(self, tensor: torch.Tensor) -> None:
|
||||
self._check_cast_functional(
|
||||
tensor, "tensor"
|
||||
).mark_mutation_hidden_from_autograd()
|
||||
|
||||
|
||||
class CppFunctionalizeAPI(BaseFunctionalizeAPI):
|
||||
def wrap_tensors(self, args: tuple[Any, ...]) -> tuple[Any, ...]:
|
||||
from torch._functorch.eager_transforms import _wrap_all_tensors_to_functional
|
||||
|
||||
return _wrap_all_tensors_to_functional(args, level=0)
|
||||
|
||||
def unwrap_tensors(
|
||||
self, args: torch.Tensor | tuple[torch.Tensor, ...]
|
||||
) -> torch.Tensor | tuple[torch.Tensor, ...]:
|
||||
from torch._functorch.eager_transforms import (
|
||||
_unwrap_all_tensors_from_functional,
|
||||
)
|
||||
|
||||
return _unwrap_all_tensors_from_functional(args, reapply_views=_reapply_views())
|
||||
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def functionalize(self, inner_f: Callable) -> Callable:
|
||||
return torch.func.functionalize(inner_f)
|
||||
|
||||
def redispatch_to_next(self) -> AbstractContextManager[None]:
|
||||
return torch._C._ExcludeDispatchKeyGuard(
|
||||
torch._C.DispatchKeySet(torch._C.DispatchKey.Functionalize)
|
||||
)
|
||||
|
||||
|
||||
class FunctorchFunctionalizeAPI(BaseFunctionalizeAPI):
|
||||
def __init__(self, interpreter: FunctionalizeInterpreter) -> None:
|
||||
self.interpreter = interpreter
|
||||
|
||||
def wrap_tensors(self, args: tuple[Any]) -> tuple[Any]:
|
||||
from torch._functorch.eager_transforms import _wrap_all_tensors_to_functional
|
||||
|
||||
return _wrap_all_tensors_to_functional(args, level=self.interpreter.level())
|
||||
|
||||
def unwrap_tensors(
|
||||
self, args: torch.Tensor | tuple[torch.Tensor, ...]
|
||||
) -> torch.Tensor | tuple[torch.Tensor, ...]:
|
||||
from torch._functorch.eager_transforms import (
|
||||
_unwrap_all_tensors_from_functional,
|
||||
)
|
||||
|
||||
return _unwrap_all_tensors_from_functional(
|
||||
args, reapply_views=self.interpreter.functionalize_add_back_views()
|
||||
)
|
||||
|
||||
# pyrefly: ignore [implicit-any]
|
||||
def functionalize(self, inner_f: Callable) -> Callable:
|
||||
return torch.func.functionalize(
|
||||
inner_f,
|
||||
remove=(
|
||||
"mutations_and_views"
|
||||
if self.interpreter.functionalize_add_back_views()
|
||||
else "mutations"
|
||||
),
|
||||
)
|
||||
|
||||
def redispatch_to_next(self) -> AbstractContextManager[None]:
|
||||
return self.interpreter.lower()
|
||||
|
||||
|
||||
def mb_unwrap_functional_tensor(tensor: torch.Tensor) -> torch.Tensor:
|
||||
if isinstance(tensor, FunctionalTensor):
|
||||
return torch._from_functional_tensor(tensor.elem)
|
||||
return tensor
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,260 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from copy import deepcopy
|
||||
from itertools import combinations
|
||||
from typing import Any, NamedTuple, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.fx.operator_schemas import _normalize_function_or_error
|
||||
from torch.utils import _pytree as pytree
|
||||
from torch.utils._python_dispatch import TorchDispatchMode
|
||||
from torch.utils._pytree import tree_map
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable
|
||||
|
||||
from torch._ops import OpOverload
|
||||
|
||||
|
||||
class Mutation(NamedTuple):
|
||||
op_name: str
|
||||
arg_name: str
|
||||
|
||||
|
||||
class Aliasing(NamedTuple):
|
||||
op_name: str
|
||||
arg_name: str
|
||||
output_number: str
|
||||
|
||||
|
||||
# Simplified naming for C++ classes
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
SchemaArgument = torch._C._SchemaArgument
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
SchemaArgType = torch._C._SchemaArgType
|
||||
SchemaInfo = torch._C._SchemaInfo
|
||||
|
||||
# This TorchDispatchMode Subclass is used to verify op schemas
|
||||
# This TorchDispatchMode Scubclass currently:
|
||||
# - Records the called ops
|
||||
# - Checks for mutations on all inputs
|
||||
# - Checks for aliasing on all inputs
|
||||
|
||||
|
||||
# move these 2 functions here to avoid numpy dependency in testing/_internal/common_utils.py
|
||||
|
||||
|
||||
def is_iterable_of_tensors(iterable: Iterable[Any]) -> bool:
|
||||
# Tensor itself is iterable so we check this first
|
||||
if isinstance(iterable, torch.Tensor):
|
||||
return False
|
||||
try:
|
||||
# pyrefly: ignore[bad-argument-type]
|
||||
if len(iterable) == 0:
|
||||
return False
|
||||
for t in iter(iterable):
|
||||
if not isinstance(t, torch.Tensor):
|
||||
return False
|
||||
except TypeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def clone_inputs(args: Iterable[Any]) -> list[Any]:
|
||||
inputs: list[Any] = []
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, torch.Tensor):
|
||||
inputs.append(arg.detach().clone())
|
||||
elif is_iterable_of_tensors(arg):
|
||||
inputs.append([t.detach().clone() for t in arg])
|
||||
else:
|
||||
inputs.append(arg)
|
||||
|
||||
return inputs
|
||||
|
||||
|
||||
class SchemaCheckMode(TorchDispatchMode):
|
||||
def __init__(self) -> None:
|
||||
# Information recorded for testing purposes. For example:
|
||||
# - incorrect schemas
|
||||
# - overly conservative schemas
|
||||
self.ops: list[str] = []
|
||||
self.mutated: list[Mutation] = []
|
||||
self.aliasing: list[Aliasing] = []
|
||||
|
||||
def reset_cache(self) -> None:
|
||||
self.ops.clear()
|
||||
self.mutated.clear()
|
||||
self.aliasing.clear()
|
||||
|
||||
def display_ops(self) -> None:
|
||||
print(*self.ops, sep=",")
|
||||
|
||||
def __torch_dispatch__(
|
||||
self,
|
||||
func: OpOverload,
|
||||
types: tuple[type[Any], ...],
|
||||
args: tuple[Any, ...] = (),
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
def bitwise_equal(lhs: torch.Tensor, rhs: torch.Tensor) -> bool:
|
||||
if lhs.is_quantized:
|
||||
# TODO: This is only OK if can't have NaN quantized; idk if
|
||||
# this is actually true
|
||||
return torch.equal(lhs, rhs)
|
||||
else:
|
||||
return torch.allclose(lhs, rhs, equal_nan=True)
|
||||
|
||||
def has_mutated(
|
||||
before: Any, after: Any, md: tuple[tuple[int, ...], int] | None
|
||||
) -> bool:
|
||||
are_tensors = type(before) is torch.Tensor and type(after) is torch.Tensor
|
||||
if (
|
||||
are_tensors
|
||||
and before.layout != torch.sparse_csr
|
||||
and after.layout != torch.sparse_csr
|
||||
):
|
||||
return md is not None and not (
|
||||
before.size() == after.size()
|
||||
and bitwise_equal(before, after)
|
||||
and md[0] == after.stride()
|
||||
and md[1] == after._typed_storage()._cdata
|
||||
)
|
||||
return False
|
||||
|
||||
def has_aliased(lhs: Any, rhs: Any) -> bool:
|
||||
try:
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return torch._C._overlaps(lhs, rhs)
|
||||
except Exception as exception:
|
||||
if str(exception).startswith("Cannot inspect value of type "):
|
||||
return False
|
||||
else:
|
||||
raise exception
|
||||
|
||||
def standardize_name(name: str) -> str:
|
||||
return name if name != "self" else "input"
|
||||
|
||||
def unwrap(e: Any) -> Any:
|
||||
if isinstance(e, torch.Tensor) and type(e) is not torch.Tensor:
|
||||
try:
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
return e.elem
|
||||
except AttributeError:
|
||||
return e
|
||||
return e
|
||||
|
||||
def parse_metadata(e: Any) -> tuple[tuple[int, ...], int] | None:
|
||||
if isinstance(e, torch.Tensor):
|
||||
if type(e) is not torch.Tensor:
|
||||
try:
|
||||
# pyrefly: ignore[missing-attribute]
|
||||
current = e.elem
|
||||
return (
|
||||
deepcopy(current.stride()),
|
||||
current._typed_storage()._cdata,
|
||||
)
|
||||
except AttributeError:
|
||||
return None
|
||||
# Sparse CSR tensors do not have strides or storage
|
||||
elif e.layout != torch.sparse_csr:
|
||||
return (deepcopy(e.stride()), e._typed_storage()._cdata)
|
||||
return None
|
||||
|
||||
self.ops.append(func._schema.name)
|
||||
|
||||
# Clone and process arguments and outputs
|
||||
pre_arguments = _normalize_function_or_error(
|
||||
func, args, kwargs, normalize_to_only_use_kwargs=True
|
||||
).kwargs
|
||||
|
||||
c_p_args = dict(zip(pre_arguments.keys(), clone_inputs(pre_arguments.values())))
|
||||
cloned_arguments = {
|
||||
name: tree_map(unwrap, c_p_args.get(name)) for name in c_p_args
|
||||
}
|
||||
cloned_metadata = {
|
||||
name: [
|
||||
parse_metadata(a) for a in pytree.tree_leaves(pre_arguments.get(name))
|
||||
]
|
||||
for name in pre_arguments
|
||||
}
|
||||
|
||||
out = func(*args, **kwargs)
|
||||
arguments = {
|
||||
name: tree_map(unwrap, pre_arguments.get(name)) for name in pre_arguments
|
||||
}
|
||||
tuple_out = out if isinstance(out, tuple) else (out,)
|
||||
tuple_out = tree_map(unwrap, tuple_out)
|
||||
|
||||
schema_info = SchemaInfo(func._schema)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
schema_info.add_argument_values(pre_arguments)
|
||||
|
||||
# Process arguments with outputs
|
||||
for i in range(len(func._schema.arguments)):
|
||||
arg = func._schema.arguments[i]
|
||||
name = standardize_name(arg.name)
|
||||
if arguments.get(name) is not None:
|
||||
before = cloned_arguments.get(name)
|
||||
md = cloned_metadata.get(name)
|
||||
after = arguments.get(name)
|
||||
for j in range(len(tuple_out)):
|
||||
# aten::_unsafe_view is intended to have incorrect aliasing notation (hence unsafe)
|
||||
unsafe_ops = ("aten::_unsafe_view", "aten::unsafe_split")
|
||||
if (
|
||||
has_aliased(tuple_out[j], after)
|
||||
and func._schema.name not in unsafe_ops
|
||||
):
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
if not schema_info.may_contain_alias(
|
||||
SchemaArgument(SchemaArgType.output, j),
|
||||
SchemaArgument(SchemaArgType.input, i),
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Argument {name} is not defined to alias output but was aliasing"
|
||||
)
|
||||
else:
|
||||
self.aliasing.append(
|
||||
Aliasing(func._schema.name, name, f"output_{j}")
|
||||
)
|
||||
if after is tuple_out[j] and isinstance(after, torch.Tensor):
|
||||
# Only mutable ops e.g. (add_, add.out) are allowed to directly return inputs.
|
||||
if not schema_info.is_mutable(
|
||||
SchemaArgument(SchemaArgType.input, i)
|
||||
) and func not in [
|
||||
torch.ops.aten.lift.default,
|
||||
torch.ops.aten.lift_fresh.default,
|
||||
]:
|
||||
raise RuntimeError(
|
||||
f"""\
|
||||
Dispatcher operators below autograd are not allowed to directly return inputs.
|
||||
However, we found that `outputs[{str(j)}] is {name}"""
|
||||
)
|
||||
if md is not None and any(
|
||||
has_mutated(a, b, c)
|
||||
for a, b, c in zip(
|
||||
pytree.tree_leaves(before), pytree.tree_leaves(after), md
|
||||
)
|
||||
):
|
||||
if not schema_info.is_mutable(
|
||||
SchemaArgument(SchemaArgType.input, i)
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Argument {name} is not defined as mutable but was mutated"
|
||||
)
|
||||
else:
|
||||
self.mutated.append(Mutation(func._schema.name, name))
|
||||
|
||||
# Aliasing between outputs
|
||||
for i, j in combinations(range(len(func._schema.returns)), 2):
|
||||
if has_aliased(tuple_out[i], tuple_out[j]):
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
if not schema_info.may_contain_alias(
|
||||
SchemaArgument(SchemaArgType.output, i),
|
||||
SchemaArgument(SchemaArgType.output, j),
|
||||
):
|
||||
raise RuntimeError(f"Outputs {i} and {j} alias unexpectedly")
|
||||
|
||||
return out
|
||||
Reference in New Issue
Block a user