Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,130 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from . import Dim
|
||||
|
||||
import torch # noqa: TC002
|
||||
|
||||
|
||||
# NB: The old code represented dimension was from as negative number, so we
|
||||
# follow this convention even though it shouldn't be necessary now
|
||||
class DimEntry:
|
||||
# The dimension this is from the rhs, or a FCD
|
||||
data: Dim | int
|
||||
|
||||
def __init__(self, data: Dim | int | None = None) -> None:
|
||||
from . import Dim
|
||||
|
||||
if type(data) is int:
|
||||
if data >= 0:
|
||||
raise AssertionError(f"Expected negative int, got {data}")
|
||||
elif data is None:
|
||||
data = 0
|
||||
else:
|
||||
if not isinstance(data, Dim):
|
||||
raise AssertionError(f"Expected Dim, got {type(data)}")
|
||||
self.data = data
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, DimEntry):
|
||||
return False
|
||||
# Use 'is' for Dim objects to avoid triggering __torch_function__
|
||||
# Use '==' only for positional (int) comparisons
|
||||
if self.is_positional() and other.is_positional():
|
||||
# Both are positional (ints)
|
||||
return self.data == other.data
|
||||
elif not self.is_positional() and not other.is_positional():
|
||||
# Both are Dim objects - use 'is' to avoid __eq__
|
||||
return self.data is other.data
|
||||
else:
|
||||
# One is positional, one is Dim - they can't be equal
|
||||
return False
|
||||
|
||||
def is_positional(self) -> bool:
|
||||
return type(self.data) is int and self.data < 0
|
||||
|
||||
def is_none(self) -> bool:
|
||||
# Use isinstance to check for Dim objects, avoid triggering __torch_function__
|
||||
from . import Dim
|
||||
|
||||
if isinstance(self.data, Dim):
|
||||
# This is a Dim object, it can't be "none" (which is represented by 0)
|
||||
return False
|
||||
else:
|
||||
# This is an int or other type
|
||||
return self.data == 0
|
||||
|
||||
def position(self) -> int:
|
||||
if not isinstance(self.data, int):
|
||||
raise AssertionError(f"Expected int, got {type(self.data)}")
|
||||
return self.data
|
||||
|
||||
def dim(self) -> Dim:
|
||||
if isinstance(self.data, int):
|
||||
raise AssertionError("Expected Dim, got int")
|
||||
return self.data
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return repr(self.data)
|
||||
|
||||
|
||||
def ndim_of_levels(levels: Sequence[DimEntry]) -> int:
|
||||
r = 0
|
||||
for l in levels:
|
||||
if l.is_positional():
|
||||
r += 1
|
||||
return r
|
||||
|
||||
|
||||
def _match_levels(
|
||||
tensor: torch.Tensor,
|
||||
from_levels: list[DimEntry],
|
||||
to_levels: list[DimEntry],
|
||||
drop_levels: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Reshape a tensor to match target levels using as_strided.
|
||||
|
||||
Args:
|
||||
tensor: Input tensor to reshape
|
||||
from_levels: Current levels of the tensor
|
||||
to_levels: Target levels to match
|
||||
drop_levels: If True, missing dimensions are assumed to have stride 0
|
||||
|
||||
Returns:
|
||||
Reshaped tensor
|
||||
"""
|
||||
if from_levels == to_levels:
|
||||
return tensor
|
||||
|
||||
sizes = tensor.size()
|
||||
strides = tensor.stride()
|
||||
|
||||
if not drop_levels:
|
||||
if len(from_levels) > len(to_levels):
|
||||
raise AssertionError("Cannot expand dimensions without drop_levels")
|
||||
|
||||
new_sizes = []
|
||||
new_strides = []
|
||||
|
||||
for level in to_levels:
|
||||
# Find index of this level in from_levels
|
||||
try:
|
||||
idx = from_levels.index(level)
|
||||
except ValueError:
|
||||
# Level not found in from_levels
|
||||
if level.is_positional():
|
||||
new_sizes.append(1)
|
||||
else:
|
||||
new_sizes.append(level.dim().size)
|
||||
new_strides.append(0)
|
||||
else:
|
||||
new_sizes.append(sizes[idx])
|
||||
new_strides.append(strides[idx])
|
||||
|
||||
return tensor.as_strided(new_sizes, new_strides, tensor.storage_offset())
|
||||
@@ -0,0 +1,143 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from ._dim_entry import DimEntry
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from . import Dim, Tensor
|
||||
|
||||
|
||||
class EnableAllLayers:
|
||||
"""
|
||||
RAII-style context manager for enabling functorch vmap layers.
|
||||
It manages the creation and cleanup of functorch dynamic layers.
|
||||
|
||||
This is probably one of the more algorithmically important parts of first
|
||||
class dims. Intuitively, FCD can be thought of as another way of using
|
||||
vmap, where you don't actually have to vmap at the top level, instead the
|
||||
vmaps are implicitly determined by inspecting the bound dimensions on the
|
||||
FCD tensors involved in a compute (this is similar to our concept of
|
||||
non-lexical modes that we spent a long time talking about years ago). But
|
||||
under the hood you still need to actually enable the vmap mode. So once
|
||||
FCD has determined all of the dims we are batching over, it needs to
|
||||
enable all those layers so functorch can actually apply the batching
|
||||
rules. Therefore enable all layers!
|
||||
"""
|
||||
|
||||
levels_start: int
|
||||
levels_to_dim: list[Dim]
|
||||
|
||||
def __init__(self, levels: list[DimEntry]):
|
||||
"""
|
||||
Initialize and push dynamic layers for all first-class dimensions.
|
||||
|
||||
Args:
|
||||
levels: List of dimension entries to create layers for
|
||||
"""
|
||||
|
||||
from . import Dim
|
||||
|
||||
self.levels_start = 0
|
||||
self.levels_to_dim = []
|
||||
|
||||
for l in levels:
|
||||
if not l.is_positional():
|
||||
d = l.dim()
|
||||
if not isinstance(d, Dim):
|
||||
raise AssertionError(f"Expected Dim, got {type(d)}")
|
||||
self.levels_to_dim.append(d)
|
||||
|
||||
# Sort by level for stable ordering
|
||||
self.levels_to_dim.sort(key=lambda d: d._level)
|
||||
|
||||
def __enter__(self) -> EnableAllLayers: # noqa: PYI034
|
||||
# Create functorch dynamic layers
|
||||
for i, dim in enumerate(self.levels_to_dim):
|
||||
batch_size = dim.size
|
||||
level = torch._C._functorch._vmap_increment_nesting(batch_size, "different")
|
||||
if i == 0:
|
||||
self.levels_start = level
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
||||
"""Clean up dynamic layers in reverse order."""
|
||||
to_remove = self.levels_start + len(self.levels_to_dim) - 1
|
||||
for i in range(len(self.levels_to_dim)):
|
||||
popped = torch._C._functorch._vmap_decrement_nesting()
|
||||
if popped != to_remove - i:
|
||||
raise AssertionError(f"Expected layer {to_remove - i}, got {popped}")
|
||||
|
||||
def from_batched(self, batchedtensor: torch.Tensor, has_device: bool) -> Tensor:
|
||||
"""
|
||||
Create a Tensor from a batched tensor by unwrapping functorch layers.
|
||||
|
||||
Args:
|
||||
batchedtensor: Batched tensor from functorch operation
|
||||
has_device: Whether tensor has device info
|
||||
|
||||
Returns:
|
||||
Tensor with appropriate levels
|
||||
"""
|
||||
# Create positional levels for base dimensions
|
||||
levels: list[DimEntry] = []
|
||||
for i in range(-batchedtensor.dim(), 0):
|
||||
levels.append(DimEntry(i))
|
||||
|
||||
tensor = batchedtensor
|
||||
|
||||
while torch._C._functorch.is_batchedtensor(tensor):
|
||||
level = torch._C._functorch.maybe_get_level(tensor)
|
||||
if level is None:
|
||||
raise AssertionError("Expected level to be non-None")
|
||||
if not (
|
||||
level >= self.levels_start
|
||||
and level < self.levels_start + len(self.levels_to_dim)
|
||||
):
|
||||
raise AssertionError(f"Level {level} out of range")
|
||||
dim = DimEntry(self.levels_to_dim[level - self.levels_start])
|
||||
bdim = torch._C._functorch.maybe_get_bdim(tensor)
|
||||
if bdim is None:
|
||||
raise AssertionError("Expected bdim to be non-None")
|
||||
levels.insert(bdim, dim)
|
||||
tensor = torch._C._functorch.get_unwrapped(tensor)
|
||||
|
||||
from . import Tensor
|
||||
|
||||
result = Tensor()
|
||||
result._tensor = tensor
|
||||
result._batchtensor = batchedtensor
|
||||
result._has_device = has_device
|
||||
result._levels = levels
|
||||
return result
|
||||
|
||||
def inplace_update_layers(
|
||||
self, batchtensor: torch.Tensor, levels: list[DimEntry]
|
||||
) -> None:
|
||||
"""
|
||||
Update the levels of a batched tensor in place.
|
||||
|
||||
This requires the _maybe_unsafe_set_level binding that we'll add to functorch.
|
||||
|
||||
Args:
|
||||
batchtensor: Batched tensor to update
|
||||
levels: New levels to set
|
||||
"""
|
||||
# Check if tensor is batched
|
||||
if not torch._C._functorch.is_batchedtensor(batchtensor):
|
||||
return
|
||||
|
||||
impl = batchtensor
|
||||
|
||||
for i in reversed(range(len(self.levels_to_dim))):
|
||||
if impl is None:
|
||||
break
|
||||
|
||||
if any(l == DimEntry(self.levels_to_dim[i]) for l in levels):
|
||||
# This is very interesting! The level on batch tensor is
|
||||
# meaningless! We set it RIGHT before we go into vmap
|
||||
torch._C._functorch._maybe_unsafe_set_level(impl, self.levels_start + i)
|
||||
impl = torch._C._functorch.get_unwrapped(impl)
|
||||
@@ -0,0 +1,564 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from ._dim_entry import _match_levels, DimEntry
|
||||
from ._tensor_info import TensorInfo
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from . import Dim
|
||||
|
||||
|
||||
def _safe_index(lst: list, item: Any) -> int | None:
|
||||
"""
|
||||
Helper function to find index of item in list.
|
||||
|
||||
For DimEntry objects, uses __eq__ comparison which properly handles
|
||||
both positional and Dim entries.
|
||||
|
||||
Returns the index if found, None if not found.
|
||||
"""
|
||||
for i, list_item in enumerate(lst):
|
||||
# Use == for DimEntry objects as they have proper __eq__ implementation
|
||||
if isinstance(item, DimEntry) and isinstance(list_item, DimEntry):
|
||||
if list_item == item:
|
||||
return i
|
||||
elif list_item is item:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexingInfo:
|
||||
can_call_original: bool = False
|
||||
advanced_indexing: bool = False
|
||||
self_tensor: torch.Tensor | None = None
|
||||
flat_inputs: list[Any] = field(default_factory=list)
|
||||
result_levels: list[DimEntry] = field(default_factory=list)
|
||||
has_device: bool = False
|
||||
|
||||
|
||||
def has_dims(obj: Any) -> bool:
|
||||
"""
|
||||
Check if an object has first-class dimensions.
|
||||
|
||||
This function checks if the object is either a Dim or a functorch Tensor
|
||||
that has first-class dimensions, using the proper check_exact methods.
|
||||
"""
|
||||
from . import Dim, Tensor
|
||||
|
||||
return Dim.check_exact(obj) or Tensor.check_exact(obj)
|
||||
|
||||
|
||||
def _bind_dims_to_size(sz: int, sd: int, dims: list, nsz: list, nsd: list) -> None:
|
||||
"""
|
||||
Bind dimensions to size and calculate proper strides for dim packs.
|
||||
"""
|
||||
from . import DimensionBindError
|
||||
|
||||
rhs_prod = 1
|
||||
for i, dim in enumerate(dims):
|
||||
if not dim.is_bound:
|
||||
# Check for multiple unbound dimensions
|
||||
for j in range(i + 1, len(dims)):
|
||||
if not dims[j].is_bound:
|
||||
raise DimensionBindError(
|
||||
f"cannot infer the sizes of two dimensions at once {dim!r} and {dims[j]!r}"
|
||||
)
|
||||
rhs_prod *= dims[j].size
|
||||
|
||||
# Calculate the size for this unbound dimension
|
||||
if sz % rhs_prod != 0:
|
||||
tup = tuple(dim.size if dim.is_bound else "?" for dim in dims)
|
||||
raise DimensionBindError(
|
||||
f"inferred dimension does not evenly fit into larger dimension: {sz} vs {tup}"
|
||||
)
|
||||
|
||||
inferred_size = sz // rhs_prod
|
||||
dim.size = inferred_size
|
||||
rhs_prod = sz
|
||||
break
|
||||
else:
|
||||
rhs_prod *= dim.size
|
||||
|
||||
# Final validation that dimensions match
|
||||
if rhs_prod != sz:
|
||||
tup = tuple(dims)
|
||||
raise DimensionBindError(
|
||||
f"Dimension sizes to do not match ({sz} != {rhs_prod}) when matching dimension pack {tup}"
|
||||
)
|
||||
|
||||
# Calculate new sizes and strides for each dimension in the pack
|
||||
# First calculate all strides by iterating in reverse
|
||||
new_strides = [0] * len(dims)
|
||||
current_stride = sd
|
||||
for i in reversed(range(len(dims))):
|
||||
new_strides[i] = current_stride
|
||||
current_stride *= dims[i].size
|
||||
|
||||
# Then append sizes and strides in forward order
|
||||
for i in range(len(dims)):
|
||||
nsz.append(dims[i].size)
|
||||
nsd.append(new_strides[i])
|
||||
|
||||
|
||||
def slice_to_tuple(flat_inputs: list) -> tuple:
|
||||
return tuple(flat_inputs)
|
||||
|
||||
|
||||
def extractIndices(index: Any, indices: list) -> bool:
|
||||
if isinstance(index, tuple): # mpy::tuple_view::check
|
||||
indices.extend(index)
|
||||
return True
|
||||
elif isinstance(index, torch.Tensor): # THPVariable_Check
|
||||
indices.append(index)
|
||||
return False
|
||||
elif not hasattr(index, "__iter__") or isinstance(
|
||||
index, (str, bytes)
|
||||
): # !mpy::is_sequence
|
||||
indices.append(index)
|
||||
return False
|
||||
|
||||
# Handle sequence case (list)
|
||||
if isinstance(index, list):
|
||||
if len(index) >= 32:
|
||||
indices.extend(index)
|
||||
return True
|
||||
|
||||
# Check each item in the sequence
|
||||
for item in index:
|
||||
if (
|
||||
isinstance(item, (torch.Tensor, slice))
|
||||
or hasattr(item, "__iter__")
|
||||
or item is ...
|
||||
or item is None
|
||||
or has_dims(item)
|
||||
):
|
||||
indices.extend(index)
|
||||
return True
|
||||
|
||||
# If we got here, treat as single index
|
||||
indices.append(index)
|
||||
return False
|
||||
|
||||
# Default case
|
||||
indices.append(index)
|
||||
return False
|
||||
|
||||
|
||||
def getitem(cls: Any, func: Any, types: Any, args: Any, kwargs: Any) -> Any:
|
||||
self = args[0]
|
||||
index = args[1]
|
||||
|
||||
iinfo = getsetitem(self, index, has_dims(self))
|
||||
if iinfo.can_call_original:
|
||||
# Call original tensor __getitem__ directly, bypassing __torch_function__
|
||||
return torch.Tensor.__getitem__(self, index)
|
||||
|
||||
return invoke_getitem(iinfo)
|
||||
|
||||
|
||||
def setitem(self: Any, index: Any, rhs: Any) -> None:
|
||||
"""Set values in tensor using first-class dimensions."""
|
||||
from . import DimensionBindError, TensorInfo
|
||||
|
||||
iinfo = getsetitem(self, index, has_dims(self) or has_dims(rhs))
|
||||
|
||||
if iinfo.can_call_original:
|
||||
# Call original tensor __setitem__ directly, bypassing __torch_function__
|
||||
torch._C.TensorBase.__setitem__(self, index, rhs)
|
||||
return
|
||||
|
||||
# Handle RHS tensor with dimensions
|
||||
rhs_info = TensorInfo.create(rhs, False, False)
|
||||
|
||||
if rhs_info:
|
||||
# Check that rhs dimensions are compatible with result dimensions
|
||||
for l in rhs_info.levels:
|
||||
if not l.is_positional():
|
||||
# Find this dimension in result levels
|
||||
found = False
|
||||
for result_level in iinfo.result_levels:
|
||||
if (
|
||||
not result_level.is_positional()
|
||||
and result_level.dim() is l.dim()
|
||||
):
|
||||
found = True
|
||||
break
|
||||
|
||||
if not found:
|
||||
# Create tuple representation of result levels for error message
|
||||
result_dims: list[int | Dim] = []
|
||||
for rl in iinfo.result_levels:
|
||||
if rl.is_positional():
|
||||
result_dims.append(rl.position())
|
||||
else:
|
||||
result_dims.append(rl.dim())
|
||||
|
||||
raise DimensionBindError(
|
||||
f"rhs of setitem contains dimension {l.dim()!r} which is not in the dimension on the left "
|
||||
f"({tuple(result_dims)!r})"
|
||||
)
|
||||
|
||||
# Match RHS tensor to result levels
|
||||
if rhs_info.tensor is None:
|
||||
raise AssertionError("Cannot match levels on None tensor")
|
||||
matched_rhs = _match_levels(
|
||||
rhs_info.tensor, rhs_info.levels, iinfo.result_levels
|
||||
)
|
||||
else:
|
||||
matched_rhs = rhs
|
||||
|
||||
# For advanced indexing with dimensions, we need special handling
|
||||
if iinfo.advanced_indexing:
|
||||
# Use advanced indexing - the flat_inputs already contain matched tensors
|
||||
tup = slice_to_tuple(iinfo.flat_inputs)
|
||||
if iinfo.self_tensor is None:
|
||||
raise RuntimeError("Cannot setitem on None tensor")
|
||||
torch._C.TensorBase.__setitem__(iinfo.self_tensor, tup, matched_rhs)
|
||||
else:
|
||||
# Simple copy operation
|
||||
if iinfo.self_tensor is None:
|
||||
raise RuntimeError("Cannot copy to None tensor")
|
||||
iinfo.self_tensor.copy_(matched_rhs)
|
||||
|
||||
|
||||
def invoke_getitem(iinfo: IndexingInfo) -> Any:
|
||||
if iinfo.advanced_indexing:
|
||||
self_tensor = iinfo.self_tensor
|
||||
tup = slice_to_tuple(iinfo.flat_inputs)
|
||||
if self_tensor is None:
|
||||
raise RuntimeError("Cannot getitem on None tensor")
|
||||
rtensor = self_tensor[tup]
|
||||
else:
|
||||
rtensor = iinfo.self_tensor # type: ignore[assignment]
|
||||
if rtensor is None:
|
||||
raise RuntimeError("Cannot getitem on None tensor")
|
||||
# rtensor is now guaranteed to be not None
|
||||
|
||||
# Create a Tensor with the proper dimensions using the class method
|
||||
from . import Tensor
|
||||
|
||||
return Tensor.from_positional(rtensor, iinfo.result_levels, iinfo.has_device)
|
||||
|
||||
|
||||
def getsetitem(self: Any, index: Any, tensors_have_dims: bool) -> IndexingInfo:
|
||||
from . import DimList # Import DimList for type checking
|
||||
|
||||
can_call_original_getitem = not tensors_have_dims
|
||||
|
||||
input_list = []
|
||||
if has_dims(index):
|
||||
input_list.append(index)
|
||||
else:
|
||||
is_sequence = extractIndices(index, input_list)
|
||||
# nothing about first class dims here, fallback to getitem
|
||||
if can_call_original_getitem and not is_sequence:
|
||||
return IndexingInfo(can_call_original=True)
|
||||
|
||||
# Calculate how many dimensions have been indexed in order to compute the
|
||||
# size of ... or expand a potentially unbound dimension list.
|
||||
dims_indexed = 0
|
||||
expanding_object = -1
|
||||
unbound_dim_list = None
|
||||
dimlists = [] # Track DimList positions for later processing
|
||||
|
||||
def check_expanding(i: int) -> None:
|
||||
nonlocal expanding_object
|
||||
if expanding_object != -1:
|
||||
from . import DimensionBindError
|
||||
|
||||
raise DimensionBindError(
|
||||
f"at most one ... or unbound dimension list can exist in indexing list but found 2 at offsets "
|
||||
f"{expanding_object} and {i}"
|
||||
)
|
||||
expanding_object = i
|
||||
|
||||
def is_dimpack(s: Any) -> bool:
|
||||
from . import Dim
|
||||
|
||||
return (
|
||||
isinstance(s, (tuple, list))
|
||||
and len(s) > 0
|
||||
and all(Dim.check_exact(item) for item in s)
|
||||
)
|
||||
|
||||
has_dimpacks_or_none = False
|
||||
for i, s in enumerate(input_list):
|
||||
if has_dims(s):
|
||||
can_call_original_getitem = False
|
||||
dims_indexed += 1
|
||||
elif s is ...:
|
||||
check_expanding(i)
|
||||
elif isinstance(s, DimList):
|
||||
can_call_original_getitem = False
|
||||
if not s.is_bound:
|
||||
check_expanding(i)
|
||||
unbound_dim_list = s
|
||||
else:
|
||||
dims_indexed += len(s._dims)
|
||||
dimlists.append(i)
|
||||
elif s is None:
|
||||
has_dimpacks_or_none = True
|
||||
elif is_dimpack(s):
|
||||
can_call_original_getitem = False
|
||||
has_dimpacks_or_none = True
|
||||
dims_indexed += 1
|
||||
else:
|
||||
dims_indexed += 1
|
||||
|
||||
# Early return if we can use original getitem
|
||||
if can_call_original_getitem:
|
||||
return IndexingInfo(can_call_original=True)
|
||||
|
||||
self_info = TensorInfo.create(self, False, True)
|
||||
total_dims = len(self_info.levels) # Total dimensions (positional + named)
|
||||
if dims_indexed > total_dims:
|
||||
raise ValueError(
|
||||
f"at least {dims_indexed} indices were supplied but the tensor only has {total_dims} dimensions"
|
||||
)
|
||||
|
||||
# Expand any unbound dimension list, or expand ... into individual : slices.
|
||||
expanding_dims = total_dims - dims_indexed
|
||||
if expanding_object != -1:
|
||||
if unbound_dim_list is not None:
|
||||
# Bind unbound dimension list to the expanding dimensions
|
||||
unbound_dim_list.bind_len(expanding_dims)
|
||||
else:
|
||||
# Expand ... into slice(None) objects
|
||||
no_slices = [slice(None)] * expanding_dims
|
||||
input_list = (
|
||||
input_list[:expanding_object]
|
||||
+ no_slices
|
||||
+ input_list[expanding_object + 1 :]
|
||||
)
|
||||
|
||||
# Flatten out any dimensions stored in dimlist elements directly into the inputs
|
||||
# Process in reverse order to maintain indices
|
||||
for i in range(len(dimlists) - 1, -1, -1):
|
||||
idx = dimlists[i]
|
||||
|
||||
# We added more elements to input because of ...
|
||||
# so we need to also adjust the index to get back to where the
|
||||
# dimlist existed
|
||||
if (
|
||||
unbound_dim_list is None
|
||||
and expanding_object != -1
|
||||
and idx > expanding_object
|
||||
):
|
||||
idx += expanding_dims
|
||||
|
||||
dl = input_list[idx]
|
||||
|
||||
# PRIVATE here naughty
|
||||
input_list = input_list[:idx] + dl._dims + input_list[idx + 1 :]
|
||||
|
||||
return getsetitem_flat(self_info, input_list, [], [], has_dimpacks_or_none)
|
||||
|
||||
|
||||
def getsetitem_flat(
|
||||
self_info: TensorInfo,
|
||||
input_list: list,
|
||||
keys: list[DimEntry],
|
||||
values: list,
|
||||
has_dimpacks_or_none: bool,
|
||||
) -> IndexingInfo:
|
||||
from . import Dim
|
||||
|
||||
# Track dimension usage
|
||||
seen_dims: list[Any] = []
|
||||
seen_dims_nuses: list[int] = []
|
||||
|
||||
def add_dim(dim: Any) -> None:
|
||||
# Use safe indexing to avoid triggering __torch_function__ on Dim objects
|
||||
idx = _safe_index(seen_dims, dim)
|
||||
if idx is not None:
|
||||
seen_dims_nuses[idx] += 1
|
||||
else:
|
||||
seen_dims.append(dim)
|
||||
seen_dims_nuses.append(1)
|
||||
|
||||
flat_inputs = []
|
||||
tensor_inputs: list[Any] = []
|
||||
device_holding_tensor = None
|
||||
|
||||
def append_flat_handle(handle: Any) -> None:
|
||||
flat_inputs.append(handle)
|
||||
tensor_inputs.append(None)
|
||||
|
||||
def append_tensor_input(ti: TensorInfo) -> None:
|
||||
flat_inputs.append(None)
|
||||
tensor_inputs.append(ti)
|
||||
nonlocal device_holding_tensor
|
||||
if ti.has_device and device_holding_tensor is None:
|
||||
device_holding_tensor = ti.tensor
|
||||
|
||||
nsz = []
|
||||
nsd = []
|
||||
if self_info.tensor is None:
|
||||
raise RuntimeError("Cannot get size/stride on None tensor")
|
||||
sz = self_info.tensor.size()
|
||||
sd = self_info.tensor.stride()
|
||||
|
||||
def append_size(i: int) -> None:
|
||||
if has_dimpacks_or_none:
|
||||
nsz.append(sz[i])
|
||||
nsd.append(sd[i])
|
||||
|
||||
input_it = input_list[:]
|
||||
|
||||
def parse_nones() -> None:
|
||||
nonlocal input_it
|
||||
while input_it and input_it[0] is None:
|
||||
append_flat_handle(slice(None))
|
||||
nsz.append(1)
|
||||
nsd.append(0)
|
||||
input_it = input_it[1:]
|
||||
|
||||
def append_item(i: int, arg: Any) -> None:
|
||||
if Dim.check_exact(arg):
|
||||
d = arg
|
||||
if d._size == -1:
|
||||
d.size = sz[i]
|
||||
add_dim(d)
|
||||
append_size(i)
|
||||
append_flat_handle(arg)
|
||||
return
|
||||
|
||||
info = TensorInfo.create(arg, False, False)
|
||||
if info:
|
||||
append_size(i)
|
||||
append_tensor_input(info)
|
||||
for level in info.levels:
|
||||
if not level.is_positional():
|
||||
add_dim(level.dim())
|
||||
return
|
||||
|
||||
if has_dimpacks_or_none:
|
||||
if isinstance(arg, (tuple, list)) and all(Dim.check_exact(d) for d in arg):
|
||||
# dim pack
|
||||
dim_pack = list(arg)
|
||||
for d in dim_pack:
|
||||
add_dim(d)
|
||||
append_flat_handle(d)
|
||||
_bind_dims_to_size(sz[i], sd[i], dim_pack, nsz, nsd)
|
||||
return
|
||||
|
||||
append_size(i)
|
||||
append_flat_handle(arg)
|
||||
|
||||
# Match indexing expressions with tensor dimensions
|
||||
for i, level in enumerate(self_info.levels):
|
||||
# Use safe indexing to avoid triggering __torch_function__ on DimEntry comparisons
|
||||
idx = _safe_index(keys, level)
|
||||
if idx is not None:
|
||||
append_item(i, values[idx])
|
||||
else:
|
||||
if level.is_positional():
|
||||
parse_nones()
|
||||
if not input_it:
|
||||
append_flat_handle(slice(None))
|
||||
append_size(i)
|
||||
else:
|
||||
arg = input_it[0]
|
||||
input_it = input_it[1:]
|
||||
append_item(i, arg)
|
||||
else:
|
||||
add_dim(level.dim())
|
||||
append_flat_handle(level.dim())
|
||||
append_size(i)
|
||||
|
||||
parse_nones()
|
||||
|
||||
# Restride tensor if needed
|
||||
if has_dimpacks_or_none and nsz:
|
||||
if self_info.tensor is None:
|
||||
raise RuntimeError("Cannot restride None tensor")
|
||||
self_tensor = self_info.tensor.as_strided(
|
||||
nsz, nsd, self_info.tensor.storage_offset()
|
||||
)
|
||||
else:
|
||||
self_tensor = self_info.tensor
|
||||
|
||||
# Determine result shape and indexing requirements
|
||||
result_levels: list[Any] = []
|
||||
index_levels = []
|
||||
tensor_insert_point = -1
|
||||
requires_getindex = False
|
||||
|
||||
def mark_tensor_index() -> None:
|
||||
nonlocal tensor_insert_point
|
||||
if tensor_insert_point == -1:
|
||||
tensor_insert_point = len(result_levels)
|
||||
elif tensor_insert_point != len(result_levels):
|
||||
tensor_insert_point = 0
|
||||
|
||||
for i, inp in enumerate(flat_inputs):
|
||||
if tensor_inputs[i] is not None:
|
||||
requires_getindex = True
|
||||
mark_tensor_index()
|
||||
for level in tensor_inputs[i].levels:
|
||||
if level not in index_levels:
|
||||
index_levels.append(level)
|
||||
elif Dim.check_exact(inp):
|
||||
d = inp
|
||||
# Use safe indexing to avoid triggering __torch_function__
|
||||
dim_idx = _safe_index(seen_dims, d)
|
||||
if dim_idx is None:
|
||||
raise AssertionError(f"Dim {d} not found in seen_dims")
|
||||
if seen_dims_nuses[dim_idx] == 1:
|
||||
flat_inputs[i] = slice(None)
|
||||
result_levels.append(DimEntry(d))
|
||||
else:
|
||||
requires_getindex = True
|
||||
flat_inputs[i] = None
|
||||
tensor_inputs[i] = TensorInfo(
|
||||
d._get_range(), [DimEntry(d)], False, None
|
||||
)
|
||||
if DimEntry(d) not in index_levels:
|
||||
index_levels.append(DimEntry(d))
|
||||
mark_tensor_index()
|
||||
else:
|
||||
if inp != slice(None):
|
||||
requires_getindex = True
|
||||
if not isinstance(inp, int):
|
||||
result_levels.append(DimEntry(-1))
|
||||
|
||||
# Insert indexing dimensions at first tensor use point
|
||||
if tensor_insert_point != -1:
|
||||
for level in reversed(index_levels):
|
||||
result_levels.insert(tensor_insert_point, level)
|
||||
|
||||
# Match tensors to indexing shape
|
||||
if requires_getindex:
|
||||
for i in range(len(flat_inputs)):
|
||||
if tensor_inputs[i] is not None:
|
||||
t = tensor_inputs[i].tensor
|
||||
if t is None:
|
||||
raise AssertionError("TensorInfo should have valid tensor data")
|
||||
if (
|
||||
not tensor_inputs[i].has_device
|
||||
and device_holding_tensor is not None
|
||||
):
|
||||
t = t.to(device_holding_tensor.device)
|
||||
flat_inputs[i] = _match_levels(t, tensor_inputs[i].levels, index_levels)
|
||||
|
||||
# Number positional dimensions correctly
|
||||
seen_positionals = 0
|
||||
for i in reversed(range(len(result_levels))):
|
||||
if result_levels[i].is_positional():
|
||||
seen_positionals += 1
|
||||
result_levels[i] = DimEntry(-seen_positionals)
|
||||
|
||||
return IndexingInfo(
|
||||
can_call_original=False,
|
||||
advanced_indexing=requires_getindex,
|
||||
self_tensor=self_tensor,
|
||||
flat_inputs=flat_inputs,
|
||||
result_levels=result_levels,
|
||||
has_device=self_info.has_device,
|
||||
)
|
||||
@@ -0,0 +1,215 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch # noqa: TC002
|
||||
|
||||
from ._dim_entry import _match_levels, DimEntry, ndim_of_levels
|
||||
|
||||
|
||||
def _wrap_dim(arg: Any, orig_ndim: int, allow_none: bool = True) -> DimEntry:
|
||||
"""
|
||||
Convert various dimension representations to DimEntry.
|
||||
|
||||
Args:
|
||||
arg: The argument to convert (Dim, int, or other)
|
||||
orig_ndim: Original number of dimensions
|
||||
allow_none: Whether to allow None values
|
||||
|
||||
Returns:
|
||||
DimEntry representation of the dimension
|
||||
"""
|
||||
from . import Dim
|
||||
|
||||
if arg is None and allow_none:
|
||||
return DimEntry() # None entry
|
||||
elif isinstance(arg, Dim):
|
||||
return DimEntry(arg)
|
||||
elif isinstance(arg, int):
|
||||
if arg < 0:
|
||||
pos = arg
|
||||
else:
|
||||
pos = arg - orig_ndim
|
||||
return DimEntry(pos)
|
||||
else:
|
||||
return DimEntry()
|
||||
|
||||
|
||||
def order(
|
||||
tensor_or_dim: torch.Tensor | Any, *dims: Any | Sequence[Any]
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Reorder the dimensions of a tensor or create a tensor from a dimension.
|
||||
|
||||
It allows reordering tensor dimensions using first-class dimensions and
|
||||
positional indices.
|
||||
|
||||
Args:
|
||||
tensor_or_dim: Input tensor with first-class dimensions, or a Dim object
|
||||
*dims: Dimensions or sequences of dimensions specifying the new order
|
||||
|
||||
Returns:
|
||||
Tensor with reordered dimensions
|
||||
|
||||
Examples:
|
||||
>>> import torch
|
||||
>>> from functorch.dim import dims
|
||||
>>> batch, channel, height, width = dims(4)
|
||||
>>> x = torch.randn(2, 3, 4, 5)[batch, channel, height, width]
|
||||
>>> # Reorder to [height, width, batch, channel]
|
||||
>>> y = order(x, height, width, batch, channel)
|
||||
"""
|
||||
from . import Dim, DimList, Tensor
|
||||
|
||||
# Handle first argument - tensor or dimension
|
||||
if isinstance(tensor_or_dim, Tensor):
|
||||
# First-class tensor
|
||||
orig_levels = tensor_or_dim._levels[:]
|
||||
data = tensor_or_dim._tensor
|
||||
has_device = tensor_or_dim._has_device
|
||||
elif isinstance(tensor_or_dim, Dim):
|
||||
# Single dimension - create range tensor
|
||||
orig_levels = [DimEntry(tensor_or_dim)]
|
||||
data = tensor_or_dim._get_range()
|
||||
has_device = False
|
||||
else:
|
||||
raise ValueError("First argument must be a Tensor or Dim object")
|
||||
|
||||
flat_positional_dims = []
|
||||
to_flatten = [] # List of (start_index, length) pairs for flattening
|
||||
levels = orig_levels[:]
|
||||
|
||||
orig_ndim = ndim_of_levels(levels)
|
||||
|
||||
def append_dim(d: DimEntry) -> None:
|
||||
"""Add a dimension to the reordering, removing it from available levels."""
|
||||
try:
|
||||
idx = levels.index(d)
|
||||
except ValueError:
|
||||
idx = None
|
||||
if idx is None:
|
||||
if d.is_positional():
|
||||
raise ValueError(
|
||||
f"tensor has {orig_ndim} positional dimensions, but {d.position() + orig_ndim} specified, "
|
||||
f"or it was specified twice"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"tensor does not contain dim {d.dim()} or it was specified twice"
|
||||
)
|
||||
|
||||
levels[idx] = DimEntry()
|
||||
flat_positional_dims.append(d)
|
||||
|
||||
n_new_positional = 0
|
||||
|
||||
# Process each dimension argument
|
||||
for arg in dims:
|
||||
entry = _wrap_dim(arg, orig_ndim, False)
|
||||
if not entry.is_none():
|
||||
append_dim(entry)
|
||||
n_new_positional += 1
|
||||
elif isinstance(arg, DimList):
|
||||
# Handle DimList
|
||||
for dim in arg._dims:
|
||||
append_dim(DimEntry(dim))
|
||||
n_new_positional += 1
|
||||
else:
|
||||
# Handle sequences of dimensions for flattening
|
||||
n_new_positional += 1
|
||||
if not hasattr(arg, "__iter__"):
|
||||
raise ValueError("expected a Dim, List[Dim], or Sequence[Dim]")
|
||||
|
||||
# Convert to list to get length
|
||||
seq = list(arg)
|
||||
to_flatten.append((len(flat_positional_dims), len(seq)))
|
||||
|
||||
for item in seq:
|
||||
entry = _wrap_dim(item, orig_ndim, False)
|
||||
if entry.is_none():
|
||||
raise ValueError("expected a Dim or int")
|
||||
append_dim(entry)
|
||||
|
||||
# Build new level ordering
|
||||
insert_point = -1
|
||||
new_levels: list[DimEntry] = []
|
||||
|
||||
# Add remaining (non-reordered) levels, finding insertion point for new dimensions
|
||||
for level in levels:
|
||||
if level.is_none():
|
||||
continue
|
||||
if level.is_positional():
|
||||
if insert_point == -1:
|
||||
insert_point = len(new_levels)
|
||||
new_levels.extend(flat_positional_dims)
|
||||
new_levels.append(level)
|
||||
|
||||
# If no positional dimensions found, append new dims at the end
|
||||
if insert_point == -1:
|
||||
insert_point = len(new_levels)
|
||||
new_levels.extend(flat_positional_dims)
|
||||
|
||||
# Match tensor to new level structure
|
||||
if data is None:
|
||||
raise AssertionError("Cannot reorder None tensor")
|
||||
ndata = _match_levels(data, orig_levels, new_levels)
|
||||
|
||||
# Handle dimension flattening if requested
|
||||
if to_flatten:
|
||||
# Now build the reshape target
|
||||
view_shape = []
|
||||
sizes = ndata.size()
|
||||
|
||||
# Add dimensions before the reordered ones
|
||||
for i in range(insert_point):
|
||||
view_shape.append(sizes[i])
|
||||
|
||||
# Process flattening groups
|
||||
i = 0
|
||||
for start_idx, length in to_flatten:
|
||||
# Add individual dims before this flattening group
|
||||
while i < start_idx:
|
||||
view_shape.append(sizes[insert_point + i])
|
||||
i += 1
|
||||
|
||||
# Flatten the group
|
||||
new_size = 1
|
||||
for j in range(length):
|
||||
new_size *= sizes[insert_point + i + j]
|
||||
view_shape.append(new_size)
|
||||
i += length
|
||||
|
||||
# Add remaining individual dims
|
||||
while i < len(flat_positional_dims):
|
||||
view_shape.append(sizes[insert_point + i])
|
||||
i += 1
|
||||
|
||||
# Add dimensions after the reordered ones
|
||||
for i in range(insert_point + len(flat_positional_dims), len(levels)):
|
||||
view_shape.append(sizes[i])
|
||||
|
||||
# Update levels by removing flattened dimensions
|
||||
n_to_remove = len(flat_positional_dims) - n_new_positional
|
||||
if n_to_remove > 0:
|
||||
# Remove flattened levels
|
||||
new_levels = (
|
||||
new_levels[:insert_point] + new_levels[insert_point + n_to_remove :]
|
||||
)
|
||||
|
||||
ndata = ndata.reshape(view_shape)
|
||||
|
||||
# Renumber positional dimensions (negative indexing from the right)
|
||||
seen = 0
|
||||
for i in range(len(new_levels) - 1, -1, -1):
|
||||
if new_levels[i].is_positional() or (
|
||||
i >= insert_point and i < insert_point + n_new_positional
|
||||
):
|
||||
seen -= 1
|
||||
new_levels[i] = DimEntry(seen)
|
||||
|
||||
result = Tensor.from_positional(ndata, new_levels, has_device)
|
||||
return result # type: ignore[return-value]
|
||||
@@ -0,0 +1,67 @@
|
||||
import dis
|
||||
from typing import Any
|
||||
|
||||
|
||||
class _PyInstDecoder:
|
||||
"""
|
||||
Decodes Python bytecode instructions to extract variable names
|
||||
"""
|
||||
|
||||
def __init__(self, code_object: Any, lasti: int) -> None:
|
||||
self.code_object = code_object
|
||||
self.instructions = list(dis.get_instructions(code_object))
|
||||
self.offset = self._find_instruction_index(lasti)
|
||||
|
||||
def _find_instruction_index(self, lasti: int) -> int:
|
||||
"""Find instruction index corresponding to lasti (byte offset)."""
|
||||
# Find the instruction at or before lasti
|
||||
# This should find the CALL instruction, not the next one
|
||||
best_idx = 0
|
||||
for i, instr in enumerate(self.instructions):
|
||||
if instr.offset <= lasti:
|
||||
best_idx = i
|
||||
else:
|
||||
break
|
||||
return best_idx
|
||||
|
||||
def next(self) -> None:
|
||||
"""Advance to the next instruction."""
|
||||
self.offset += 1
|
||||
|
||||
def opcode(self) -> str | None:
|
||||
"""Get the opcode name of the current instruction."""
|
||||
if self.offset < len(self.instructions):
|
||||
return self.instructions[self.offset].opname
|
||||
return None
|
||||
|
||||
def oparg(self) -> int:
|
||||
"""Get the argument of the current instruction."""
|
||||
if self.offset < len(self.instructions):
|
||||
return self.instructions[self.offset].arg or 0
|
||||
return 0
|
||||
|
||||
def name(self) -> str | None:
|
||||
"""
|
||||
Extract variable name from current instruction.
|
||||
"""
|
||||
opname = self.opcode()
|
||||
if not opname:
|
||||
return None
|
||||
|
||||
names = None
|
||||
if opname in ("STORE_NAME", "STORE_GLOBAL"):
|
||||
names = self.code_object.co_names
|
||||
elif opname == "STORE_FAST":
|
||||
names = self.code_object.co_varnames
|
||||
elif opname == "STORE_DEREF":
|
||||
names = self.code_object.co_cellvars
|
||||
if not names:
|
||||
names = self.code_object.co_freevars
|
||||
else:
|
||||
return None
|
||||
|
||||
arg = self.oparg()
|
||||
if names and 0 <= arg < len(names):
|
||||
return names[arg]
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,69 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._dim_entry import DimEntry
|
||||
|
||||
|
||||
@dataclass
|
||||
class TensorInfo:
|
||||
tensor: torch.Tensor | None
|
||||
levels: list[DimEntry]
|
||||
has_device: bool
|
||||
batchedtensor: torch.Tensor | None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
from ._dim_entry import DimEntry
|
||||
|
||||
if not all(isinstance(l, DimEntry) for l in self.levels):
|
||||
raise AssertionError("All levels must be DimEntry instances")
|
||||
|
||||
def ndim(self) -> int:
|
||||
from ._dim_entry import ndim_of_levels
|
||||
|
||||
return ndim_of_levels(self.levels)
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return self.tensor is not None
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
h: Any, ensure_batched: bool = True, ensure_present: bool = True
|
||||
) -> TensorInfo:
|
||||
from . import Dim, DimEntry, Tensor
|
||||
|
||||
if Tensor.check_exact(h):
|
||||
# functorch Tensor with first-class dimensions
|
||||
return TensorInfo(
|
||||
h._get_tensor(),
|
||||
h._get_levels(),
|
||||
h._get_has_device(),
|
||||
h._get_batchtensor() if ensure_batched else None,
|
||||
)
|
||||
elif Dim.check_exact(h):
|
||||
# For Dim objects, only get range/batchtensor if needed and dimension is bound
|
||||
tensor = h._get_range() if h.is_bound else None
|
||||
batchtensor = (
|
||||
h._get_batchtensor() if ensure_batched and h.is_bound else None
|
||||
)
|
||||
return TensorInfo(
|
||||
tensor,
|
||||
[DimEntry(h)],
|
||||
False,
|
||||
batchtensor,
|
||||
)
|
||||
elif isinstance(h, torch.Tensor):
|
||||
# Plain torch tensor - create positional levels
|
||||
levels = []
|
||||
for i in range(-h.dim(), 0):
|
||||
levels.append(DimEntry(i))
|
||||
return TensorInfo(h, levels, True, h)
|
||||
else:
|
||||
if ensure_present:
|
||||
raise ValueError("expected a tensor object")
|
||||
return TensorInfo(None, [], False, None)
|
||||
@@ -0,0 +1,269 @@
|
||||
"""
|
||||
Python implementation of function wrapping functionality for functorch.dim.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch.utils._pytree import tree_map
|
||||
|
||||
from ._dim_entry import DimEntry
|
||||
from ._enable_all_layers import EnableAllLayers
|
||||
from ._tensor_info import TensorInfo
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
|
||||
def handle_from_tensor(tensor: torch.Tensor) -> torch.Tensor:
|
||||
"""Handle tensor conversion for torch function integration."""
|
||||
return tensor
|
||||
|
||||
|
||||
class WrappedOperator:
|
||||
"""
|
||||
This class wraps PyTorch operations to support first-class dimensions.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, orig: Callable, wrapper_implementation: Callable, dim_name: str = "dim"
|
||||
):
|
||||
self.orig = orig
|
||||
self.wrapper_implementation = wrapper_implementation
|
||||
self.name = getattr(orig, "__name__", "")
|
||||
self.doc = getattr(orig, "__doc__", None)
|
||||
self.dim_name = dim_name
|
||||
|
||||
self.is_pointwise = False
|
||||
self.dim_offset = 0
|
||||
self.keepdim_offset = 1
|
||||
self.single_dim = False
|
||||
self.reduce = True
|
||||
|
||||
# Update docstring if we have a dim_name
|
||||
if self.doc and self.dim_name:
|
||||
self.doc = f"{self.doc}\nArgument '{self.dim_name}' can be either an integer or a torchdim.Dim object.\n"
|
||||
|
||||
def function(self) -> Callable:
|
||||
"""Create a wrapped function that calls our wrapper implementation."""
|
||||
|
||||
def wrapped_func(*args: Any, **kwargs: Any) -> Any:
|
||||
return self.wrapper_implementation(self, *args, **kwargs)
|
||||
|
||||
# Copy metadata using functools.update_wrapper for just __name__ and __doc__
|
||||
functools.update_wrapper(
|
||||
wrapped_func, self.orig, assigned=("__name__",), updated=()
|
||||
)
|
||||
wrapped_func.__doc__ = self.doc
|
||||
|
||||
return wrapped_func
|
||||
|
||||
|
||||
def _wrap_dim(dim: Any, ndim: int, keepdim: bool = False) -> DimEntry:
|
||||
"""Convert single dimension specification to DimEntry object."""
|
||||
from . import Dim
|
||||
|
||||
if isinstance(dim, Dim):
|
||||
if keepdim:
|
||||
raise ValueError("cannot preserve first-class dimensions with keepdim=True")
|
||||
return DimEntry(dim)
|
||||
elif isinstance(dim, int):
|
||||
i = dim
|
||||
while i >= 0:
|
||||
i -= ndim
|
||||
return DimEntry(i)
|
||||
else:
|
||||
return DimEntry()
|
||||
|
||||
|
||||
def _wrap_dims(dim: Any, ndim: int, keepdim: bool = False) -> list[DimEntry]:
|
||||
"""Convert dimension specification to list of DimEntry objects."""
|
||||
de = _wrap_dim(dim, ndim, keepdim)
|
||||
result = []
|
||||
if not de.is_none():
|
||||
result.append(de)
|
||||
else:
|
||||
for d in dim:
|
||||
result.append(_wrap_dim(d, ndim, keepdim))
|
||||
return result
|
||||
|
||||
|
||||
def patched_dim_method(wrapper: WrappedOperator, *args: Any, **kwargs: Any) -> Any:
|
||||
"""
|
||||
This is the core method that handles dimension-aware operations.
|
||||
"""
|
||||
if not args:
|
||||
raise ValueError("Expected at least one argument (self)")
|
||||
|
||||
# Get dimension argument
|
||||
dim_arg = kwargs.get(wrapper.dim_name)
|
||||
if dim_arg is None and wrapper.dim_offset < len(args):
|
||||
# Try to get dim from positional args (accounting for self at index 0)
|
||||
dim_idx = wrapper.dim_offset + 1
|
||||
if dim_idx < len(args):
|
||||
dim_arg = args[dim_idx]
|
||||
|
||||
# If no dimension argument provided, fall back to standard functorch handling
|
||||
if dim_arg is None:
|
||||
info = TensorInfo.create(args[0], ensure_batched=True, ensure_present=False)
|
||||
if not info:
|
||||
return wrapper.orig(*args, **kwargs)
|
||||
|
||||
with EnableAllLayers(info.levels) as guard:
|
||||
if info.batchedtensor is None:
|
||||
raise AssertionError("Expected batchedtensor to be non-None")
|
||||
guard.inplace_update_layers(info.batchedtensor, info.levels)
|
||||
new_args = list(args)
|
||||
new_args[0] = handle_from_tensor(info.batchedtensor)
|
||||
result = wrapper.orig(*new_args, **kwargs)
|
||||
return guard.from_batched(result, info.has_device)
|
||||
|
||||
# Handle dimension-aware operation
|
||||
info = TensorInfo.create(args[0])
|
||||
if not info:
|
||||
return wrapper.orig(*args, **kwargs)
|
||||
|
||||
# Check for keepdim parameter
|
||||
keepdim = False
|
||||
if wrapper.reduce:
|
||||
keepdim_arg = kwargs.get("keepdim")
|
||||
if keepdim_arg is None and wrapper.keepdim_offset < len(args):
|
||||
keepdim_idx = wrapper.keepdim_offset + 1
|
||||
if keepdim_idx < len(args):
|
||||
keepdim_arg = args[keepdim_idx]
|
||||
if keepdim_arg is not None:
|
||||
keepdim = bool(keepdim_arg)
|
||||
|
||||
# Wrap dimensions
|
||||
ndim = info.ndim()
|
||||
dims = _wrap_dims(dim_arg, ndim, keepdim)
|
||||
|
||||
# Convert dimensions to indices and validate
|
||||
dim_indices: list[int] = []
|
||||
seen = [False] * len(info.levels)
|
||||
|
||||
for d in dims:
|
||||
midx = None
|
||||
for i, level in enumerate(info.levels):
|
||||
if level == d:
|
||||
midx = i
|
||||
break
|
||||
|
||||
if midx is None:
|
||||
# Try to match by position/name more flexibly
|
||||
for i, level in enumerate(info.levels):
|
||||
if hasattr(level, "matches") and level.matches(d):
|
||||
midx = i
|
||||
break
|
||||
|
||||
if midx is None:
|
||||
level_strs = [str(level) for level in info.levels]
|
||||
raise ValueError(
|
||||
f"Tensor with dimensions {level_strs} does not contain {d}"
|
||||
)
|
||||
|
||||
seen[midx] = True
|
||||
dim_indices.append(midx)
|
||||
|
||||
# Determine new levels after reduction
|
||||
new_levels = []
|
||||
if wrapper.reduce and not keepdim:
|
||||
for i, level in enumerate(info.levels):
|
||||
if not seen[i]:
|
||||
new_levels.append(level)
|
||||
else:
|
||||
new_levels = info.levels[:]
|
||||
|
||||
# Create dimension indices for the original function
|
||||
if len(dim_indices) == 1:
|
||||
py_indices: Any = dim_indices[0]
|
||||
else:
|
||||
py_indices = tuple(dim_indices)
|
||||
|
||||
# Update arguments
|
||||
new_args = list(args)
|
||||
new_kwargs = kwargs.copy()
|
||||
if info.tensor is None:
|
||||
raise AssertionError("Expected tensor to be non-None")
|
||||
new_args[0] = handle_from_tensor(info.tensor)
|
||||
|
||||
# Update dimension argument
|
||||
if wrapper.dim_name in new_kwargs:
|
||||
new_kwargs[wrapper.dim_name] = py_indices
|
||||
else:
|
||||
dim_idx = wrapper.dim_offset + 1
|
||||
if dim_idx < len(new_args):
|
||||
new_args = list(new_args)
|
||||
new_args[dim_idx] = py_indices
|
||||
|
||||
# Call original function
|
||||
result = wrapper.orig(*new_args, **new_kwargs)
|
||||
|
||||
# Wrap results
|
||||
def wrap_result(obj: Any) -> Any:
|
||||
if isinstance(obj, torch.Tensor):
|
||||
from . import Tensor
|
||||
|
||||
return Tensor.from_positional(obj, new_levels, info.has_device)
|
||||
return obj
|
||||
|
||||
return tree_map(wrap_result, result)
|
||||
|
||||
|
||||
def _wrap(
|
||||
orig: Callable,
|
||||
dim_offset: int | None = None,
|
||||
keepdim_offset: int | None = None,
|
||||
dim_name: str | None = None,
|
||||
single_dim: bool | None = None,
|
||||
reduce: bool | None = None,
|
||||
) -> Callable:
|
||||
"""
|
||||
Wrap a PyTorch function to support first-class dimensions.
|
||||
|
||||
Args:
|
||||
orig: Original function to wrap
|
||||
dim_offset: Offset for dimension argument (default: 0)
|
||||
keepdim_offset: Offset for keepdim argument (default: 1)
|
||||
dim_name: Name of dimension parameter (default: "dim")
|
||||
single_dim: Whether function takes single dimension (default: False)
|
||||
reduce: Whether function reduces dimensions (default: True)
|
||||
"""
|
||||
dim_name = dim_name or "dim"
|
||||
|
||||
wrapper = WrappedOperator(orig, patched_dim_method, dim_name)
|
||||
|
||||
if dim_offset is not None:
|
||||
wrapper.dim_offset = dim_offset
|
||||
if keepdim_offset is not None:
|
||||
wrapper.keepdim_offset = keepdim_offset
|
||||
if single_dim is not None:
|
||||
wrapper.single_dim = single_dim
|
||||
if reduce is not None:
|
||||
wrapper.reduce = reduce
|
||||
|
||||
return wrapper.function()
|
||||
|
||||
|
||||
def call_torch_function(
|
||||
wrapper: WrappedOperator,
|
||||
func: Callable,
|
||||
types: tuple,
|
||||
args: tuple = (),
|
||||
kwargs: dict | None = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Handle __torch_function__ calls for wrapped operators.
|
||||
"""
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from . import _Tensor
|
||||
|
||||
# Use the torch function mechanism from _Tensor
|
||||
return _Tensor.__torch_function__(func, types, args, kwargs)
|
||||
@@ -0,0 +1,48 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD-style license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
|
||||
|
||||
@contextmanager
|
||||
def magic_trace(
|
||||
output: str = "trace.fxt", magic_trace_cache: str = "/tmp/magic-trace"
|
||||
) -> Generator[None, None, None]:
|
||||
pid = os.getpid()
|
||||
if not os.path.exists(magic_trace_cache):
|
||||
print(f"Downloading magic_trace to: {magic_trace_cache}")
|
||||
subprocess.run(
|
||||
[
|
||||
"wget",
|
||||
"-O",
|
||||
magic_trace_cache,
|
||||
"-q",
|
||||
"https://github.com/janestreet/magic-trace/releases/download/v1.0.2/magic-trace",
|
||||
]
|
||||
)
|
||||
subprocess.run(["chmod", "+x", magic_trace_cache])
|
||||
args = [magic_trace_cache, "attach", "-pid", str(pid), "-o", output]
|
||||
p = subprocess.Popen(args, stderr=subprocess.PIPE, encoding="utf-8")
|
||||
if p.stderr is None:
|
||||
raise AssertionError("Expected stderr to be non-None")
|
||||
while True:
|
||||
x = p.stderr.readline()
|
||||
print(x)
|
||||
if "Attached" in x:
|
||||
break
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
p.send_signal(signal.SIGINT)
|
||||
r = p.wait()
|
||||
if p.stderr is not None:
|
||||
print(p.stderr.read())
|
||||
p.stderr.close()
|
||||
if r != 0:
|
||||
raise ValueError(f"magic_trace exited abnormally: {r}")
|
||||
@@ -0,0 +1,312 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD-style license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import torch
|
||||
|
||||
|
||||
# pointwise operators can go through a faster pathway
|
||||
|
||||
tensor_magic_methods = ["add", ""]
|
||||
pointwise_magic_methods_with_reverse = (
|
||||
"add",
|
||||
"sub",
|
||||
"mul",
|
||||
"floordiv",
|
||||
"div",
|
||||
"truediv",
|
||||
"mod",
|
||||
"pow",
|
||||
"lshift",
|
||||
"rshift",
|
||||
"and",
|
||||
"or",
|
||||
"xor",
|
||||
)
|
||||
pointwise_magic_methods = (
|
||||
*(x for m in pointwise_magic_methods_with_reverse for x in (m, "r" + m)),
|
||||
"eq",
|
||||
"gt",
|
||||
"le",
|
||||
"lt",
|
||||
"ge",
|
||||
"gt",
|
||||
"ne",
|
||||
"neg",
|
||||
"pos",
|
||||
"abs",
|
||||
"invert",
|
||||
"iadd",
|
||||
"isub",
|
||||
"imul",
|
||||
"ifloordiv",
|
||||
"idiv",
|
||||
"itruediv",
|
||||
"imod",
|
||||
"ipow",
|
||||
"ilshift",
|
||||
"irshift",
|
||||
"iand",
|
||||
"ior",
|
||||
"ixor",
|
||||
"int",
|
||||
"long",
|
||||
"float",
|
||||
"complex",
|
||||
)
|
||||
|
||||
pointwise_methods = (*(f"__{m}__" for m in pointwise_magic_methods),)
|
||||
|
||||
pointwise = (
|
||||
*(getattr(torch.Tensor, m) for m in pointwise_methods),
|
||||
torch.nn.functional.dropout,
|
||||
torch.where,
|
||||
torch.Tensor.abs,
|
||||
torch.abs,
|
||||
torch.Tensor.acos,
|
||||
torch.acos,
|
||||
torch.Tensor.acosh,
|
||||
torch.acosh,
|
||||
torch.Tensor.add,
|
||||
torch.add,
|
||||
torch.Tensor.addcdiv,
|
||||
torch.addcdiv,
|
||||
torch.Tensor.addcmul,
|
||||
torch.addcmul,
|
||||
torch.Tensor.addr,
|
||||
torch.addr,
|
||||
torch.Tensor.angle,
|
||||
torch.angle,
|
||||
torch.Tensor.asin,
|
||||
torch.asin,
|
||||
torch.Tensor.asinh,
|
||||
torch.asinh,
|
||||
torch.Tensor.atan,
|
||||
torch.atan,
|
||||
torch.Tensor.atan2,
|
||||
torch.atan2,
|
||||
torch.Tensor.atanh,
|
||||
torch.atanh,
|
||||
torch.Tensor.bitwise_and,
|
||||
torch.bitwise_and,
|
||||
torch.Tensor.bitwise_left_shift,
|
||||
torch.bitwise_left_shift,
|
||||
torch.Tensor.bitwise_not,
|
||||
torch.bitwise_not,
|
||||
torch.Tensor.bitwise_or,
|
||||
torch.bitwise_or,
|
||||
torch.Tensor.bitwise_right_shift,
|
||||
torch.bitwise_right_shift,
|
||||
torch.Tensor.bitwise_xor,
|
||||
torch.bitwise_xor,
|
||||
torch.Tensor.ceil,
|
||||
torch.ceil,
|
||||
torch.celu,
|
||||
torch.nn.functional.celu,
|
||||
torch.Tensor.clamp,
|
||||
torch.clamp,
|
||||
torch.Tensor.clamp_max,
|
||||
torch.clamp_max,
|
||||
torch.Tensor.clamp_min,
|
||||
torch.clamp_min,
|
||||
torch.Tensor.copysign,
|
||||
torch.copysign,
|
||||
torch.Tensor.cos,
|
||||
torch.cos,
|
||||
torch.Tensor.cosh,
|
||||
torch.cosh,
|
||||
torch.Tensor.deg2rad,
|
||||
torch.deg2rad,
|
||||
torch.Tensor.digamma,
|
||||
torch.digamma,
|
||||
torch.Tensor.div,
|
||||
torch.div,
|
||||
torch.dropout,
|
||||
torch.nn.functional.dropout,
|
||||
torch.nn.functional.elu,
|
||||
torch.Tensor.eq,
|
||||
torch.eq,
|
||||
torch.Tensor.erf,
|
||||
torch.erf,
|
||||
torch.Tensor.erfc,
|
||||
torch.erfc,
|
||||
torch.Tensor.erfinv,
|
||||
torch.erfinv,
|
||||
torch.Tensor.exp,
|
||||
torch.exp,
|
||||
torch.Tensor.exp2,
|
||||
torch.exp2,
|
||||
torch.Tensor.expm1,
|
||||
torch.expm1,
|
||||
torch.feature_dropout,
|
||||
torch.Tensor.float_power,
|
||||
torch.float_power,
|
||||
torch.Tensor.floor,
|
||||
torch.floor,
|
||||
torch.Tensor.floor_divide,
|
||||
torch.floor_divide,
|
||||
torch.Tensor.fmod,
|
||||
torch.fmod,
|
||||
torch.Tensor.frac,
|
||||
torch.frac,
|
||||
torch.Tensor.frexp,
|
||||
torch.frexp,
|
||||
torch.Tensor.gcd,
|
||||
torch.gcd,
|
||||
torch.Tensor.ge,
|
||||
torch.ge,
|
||||
torch.nn.functional.gelu,
|
||||
torch.nn.functional.glu,
|
||||
torch.Tensor.gt,
|
||||
torch.gt,
|
||||
torch.Tensor.hardshrink,
|
||||
torch.hardshrink,
|
||||
torch.nn.functional.hardshrink,
|
||||
torch.nn.functional.hardsigmoid,
|
||||
torch.nn.functional.hardswish,
|
||||
torch.nn.functional.hardtanh,
|
||||
torch.Tensor.heaviside,
|
||||
torch.heaviside,
|
||||
torch.Tensor.hypot,
|
||||
torch.hypot,
|
||||
torch.Tensor.i0,
|
||||
torch.i0,
|
||||
torch.Tensor.igamma,
|
||||
torch.igamma,
|
||||
torch.Tensor.igammac,
|
||||
torch.igammac,
|
||||
torch.Tensor.isclose,
|
||||
torch.isclose,
|
||||
torch.Tensor.isfinite,
|
||||
torch.isfinite,
|
||||
torch.Tensor.isinf,
|
||||
torch.isinf,
|
||||
torch.Tensor.isnan,
|
||||
torch.isnan,
|
||||
torch.Tensor.isneginf,
|
||||
torch.isneginf,
|
||||
torch.Tensor.isposinf,
|
||||
torch.isposinf,
|
||||
torch.Tensor.isreal,
|
||||
torch.isreal,
|
||||
torch.Tensor.kron,
|
||||
torch.kron,
|
||||
torch.Tensor.lcm,
|
||||
torch.lcm,
|
||||
torch.Tensor.ldexp,
|
||||
torch.ldexp,
|
||||
torch.Tensor.le,
|
||||
torch.le,
|
||||
torch.nn.functional.leaky_relu,
|
||||
torch.Tensor.lerp,
|
||||
torch.lerp,
|
||||
torch.Tensor.lgamma,
|
||||
torch.lgamma,
|
||||
torch.Tensor.log,
|
||||
torch.log,
|
||||
torch.Tensor.log10,
|
||||
torch.log10,
|
||||
torch.Tensor.log1p,
|
||||
torch.log1p,
|
||||
torch.Tensor.log2,
|
||||
torch.log2,
|
||||
torch.nn.functional.logsigmoid,
|
||||
torch.Tensor.logical_and,
|
||||
torch.logical_and,
|
||||
torch.Tensor.logical_not,
|
||||
torch.logical_not,
|
||||
torch.Tensor.logical_or,
|
||||
torch.logical_or,
|
||||
torch.Tensor.logical_xor,
|
||||
torch.logical_xor,
|
||||
torch.Tensor.logit,
|
||||
torch.logit,
|
||||
torch.Tensor.lt,
|
||||
torch.lt,
|
||||
torch.Tensor.maximum,
|
||||
torch.maximum,
|
||||
torch.Tensor.minimum,
|
||||
torch.minimum,
|
||||
torch.nn.functional.mish,
|
||||
torch.Tensor.mvlgamma,
|
||||
torch.mvlgamma,
|
||||
torch.Tensor.nan_to_num,
|
||||
torch.nan_to_num,
|
||||
torch.Tensor.ne,
|
||||
torch.ne,
|
||||
torch.Tensor.neg,
|
||||
torch.neg,
|
||||
torch.Tensor.nextafter,
|
||||
torch.nextafter,
|
||||
torch.Tensor.outer,
|
||||
torch.outer,
|
||||
torch.polar,
|
||||
torch.Tensor.polygamma,
|
||||
torch.polygamma,
|
||||
torch.Tensor.positive,
|
||||
torch.positive,
|
||||
torch.Tensor.pow,
|
||||
torch.pow,
|
||||
torch.Tensor.prelu,
|
||||
torch.prelu,
|
||||
torch.nn.functional.prelu,
|
||||
torch.Tensor.rad2deg,
|
||||
torch.rad2deg,
|
||||
torch.Tensor.reciprocal,
|
||||
torch.reciprocal,
|
||||
torch.Tensor.relu,
|
||||
torch.relu,
|
||||
torch.nn.functional.relu,
|
||||
torch.nn.functional.relu6,
|
||||
torch.Tensor.remainder,
|
||||
torch.remainder,
|
||||
torch.Tensor.round,
|
||||
torch.round,
|
||||
torch.rrelu,
|
||||
torch.nn.functional.rrelu,
|
||||
torch.Tensor.rsqrt,
|
||||
torch.rsqrt,
|
||||
torch.rsub,
|
||||
torch.selu,
|
||||
torch.nn.functional.selu,
|
||||
torch.Tensor.sgn,
|
||||
torch.sgn,
|
||||
torch.Tensor.sigmoid,
|
||||
torch.sigmoid,
|
||||
torch.nn.functional.sigmoid,
|
||||
torch.Tensor.sign,
|
||||
torch.sign,
|
||||
torch.Tensor.signbit,
|
||||
torch.signbit,
|
||||
torch.nn.functional.silu,
|
||||
torch.Tensor.sin,
|
||||
torch.sin,
|
||||
torch.Tensor.sinc,
|
||||
torch.sinc,
|
||||
torch.Tensor.sinh,
|
||||
torch.sinh,
|
||||
torch.nn.functional.softplus,
|
||||
torch.nn.functional.softshrink,
|
||||
torch.Tensor.sqrt,
|
||||
torch.sqrt,
|
||||
torch.Tensor.square,
|
||||
torch.square,
|
||||
torch.Tensor.sub,
|
||||
torch.sub,
|
||||
torch.Tensor.tan,
|
||||
torch.tan,
|
||||
torch.Tensor.tanh,
|
||||
torch.tanh,
|
||||
torch.nn.functional.tanh,
|
||||
torch.threshold,
|
||||
torch.nn.functional.threshold,
|
||||
torch.trapz,
|
||||
torch.Tensor.true_divide,
|
||||
torch.true_divide,
|
||||
torch.Tensor.trunc,
|
||||
torch.trunc,
|
||||
torch.Tensor.xlogy,
|
||||
torch.xlogy,
|
||||
torch.rand_like,
|
||||
)
|
||||
@@ -0,0 +1,71 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the BSD-style license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
from types import (
|
||||
BuiltinMethodType,
|
||||
FunctionType,
|
||||
GetSetDescriptorType,
|
||||
MethodDescriptorType,
|
||||
WrapperDescriptorType,
|
||||
)
|
||||
from typing import Any
|
||||
|
||||
|
||||
FUNC_TYPES = (
|
||||
FunctionType,
|
||||
MethodDescriptorType,
|
||||
BuiltinMethodType,
|
||||
WrapperDescriptorType,
|
||||
)
|
||||
PROPERTY_TYPES = (GetSetDescriptorType, property)
|
||||
|
||||
|
||||
def _py_wrap_method(orig: Callable, __torch_function__: Callable) -> Callable:
|
||||
def impl(*args: Any, **kwargs: Any) -> Any:
|
||||
return __torch_function__(orig, None, args, kwargs)
|
||||
|
||||
# Copy metadata using functools.update_wrapper for just __name__ and __doc__
|
||||
functools.update_wrapper(impl, orig, assigned=("__name__", "__doc__"), updated=())
|
||||
|
||||
return impl
|
||||
|
||||
|
||||
def wrap_type(to_patch: Any, pattern: type, __torch_function__: Callable) -> None:
|
||||
wrap_method = _py_wrap_method
|
||||
|
||||
all: dict[str, Any] = {}
|
||||
for t in reversed(pattern.mro()[:-1]): # skip object
|
||||
all.update(t.__dict__)
|
||||
|
||||
def wrap_attr(orig: Any) -> property:
|
||||
return property(wrap_method(orig.__get__, __torch_function__))
|
||||
|
||||
for name, obj in all.items():
|
||||
if name in (
|
||||
"__dict__",
|
||||
"__new__",
|
||||
"__init__",
|
||||
"__repr__",
|
||||
"__weakref__",
|
||||
"__doc__",
|
||||
"__module__",
|
||||
"__dir__",
|
||||
):
|
||||
continue
|
||||
|
||||
# skip things that have been overloaded
|
||||
# things that come from object like `__eq__` still need to be patched, however.
|
||||
if hasattr(to_patch, name) and getattr(to_patch, name) is not getattr(
|
||||
object, name, None
|
||||
):
|
||||
continue
|
||||
|
||||
if isinstance(obj, FUNC_TYPES):
|
||||
setattr(to_patch, name, wrap_method(obj, __torch_function__))
|
||||
elif isinstance(obj, PROPERTY_TYPES):
|
||||
setattr(to_patch, name, wrap_attr(obj))
|
||||
Reference in New Issue
Block a user