| # 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. |
| |
| # reference python implementations for C ops |
| import torch |
| from functorch._C import dim as _C |
| |
| from . import op_properties |
| from .batch_tensor import _enable_layers |
| from .tree_map import tree_flatten, tree_map |
| |
| |
| DimList = _C.DimList |
| import operator |
| from functools import reduce |
| |
| |
| # use dict to avoid writing C++ bindings for set |
| pointwise = set(op_properties.pointwise) |
| |
| |
| def prod(x): |
| return reduce(operator.mul, x, 1) |
| |
| |
| def _wrap_dim(d, N, keepdim): |
| from . import Dim |
| |
| if isinstance(d, Dim): |
| assert not keepdim, "cannot preserve first-class dimensions with keepdim=True" |
| return d |
| elif d >= 0: |
| return d - N |
| else: |
| return d |
| |
| |
| def _dims(d, N, keepdim, single_dim): |
| from . import Dim |
| |
| if isinstance(d, (Dim, int)): |
| return ltuple((_wrap_dim(d, N, keepdim),)) |
| assert not single_dim, f"expected a single dimension or int but found: {d}" |
| return ltuple(_wrap_dim(x, N, keepdim) for x in d) |
| |
| |
| def _bind_dims_to_size(lhs_size, rhs, lhs_debug): |
| from . import DimensionMismatchError |
| |
| not_bound = tuple((i, r) for i, r in enumerate(rhs) if not r.is_bound) |
| if len(not_bound) == 1: |
| idx, d = not_bound[0] |
| rhs_so_far = prod(r.size for r in rhs if r.is_bound) |
| if lhs_size % rhs_so_far != 0: |
| rhs_s = tuple("?" if not r.is_bound else str(r.size) for r in rhs) |
| raise DimensionMismatchError( |
| f"inferred dimension does not evenly fit into larger dimension: {lhs_size} vs {rhs_s}" |
| ) |
| new_size = lhs_size // rhs_so_far |
| d.size = new_size |
| elif len(not_bound) > 1: |
| rhs_s = tuple("?" if not r.is_bound else str(r.size) for r in rhs) |
| raise DimensionMismatchError( |
| f"cannot infer the size of two dimensions at once: {rhs} with sizes {rhs_s}" |
| ) |
| else: |
| rhs_size = prod(r.size for r in rhs) |
| if lhs_size != rhs_size: |
| raise DimensionMismatchError( |
| f"Dimension sizes to do not match ({lhs_size} != {rhs_size}) when matching {lhs_debug} to {rhs}" |
| ) |
| |
| |
| def _tensor_levels(inp): |
| from . import _Tensor |
| |
| if isinstance(inp, _Tensor): |
| return inp._tensor, llist(inp._levels), inp._has_device |
| else: |
| return inp, llist(range(-inp.ndim, 0)), True |
| |
| |
| def _match_levels(v, from_levels, to_levels): |
| view = [] |
| permute = [] |
| requires_view = False |
| size = v.size() |
| for t in to_levels: |
| try: |
| idx = from_levels.index(t) |
| permute.append(idx) |
| view.append(size[idx]) |
| except ValueError: |
| view.append(1) |
| requires_view = True |
| if permute != list(range(len(permute))): |
| v = v.permute(*permute) |
| if requires_view: |
| v = v.view(*view) |
| return v |
| |
| |
| # make a single dimension positional but do not permute it, |
| # used to do multi-tensor operators where the dim being acted on |
| # should not physically move if possible |
| def _positional_no_permute(self, dim, expand_dim=False): |
| from . import Tensor |
| |
| ptensor, levels = self._tensor, llist(self._levels) |
| try: |
| idx = levels.index(dim) |
| except ValueError: |
| if not expand_dim: |
| raise |
| idx = 0 |
| ptensor = ptensor.expand(dim.size, *ptensor.size()) |
| levels.insert(0, 0) |
| idx_batched = 0 |
| for i in range(idx): |
| if isinstance(levels[i], int): |
| levels[i] -= 1 |
| idx_batched += 1 |
| levels[idx] = -idx_batched - 1 |
| return Tensor.from_positional(ptensor, levels, self._has_device), idx_batched |
| |
| |
| def seq(a, b): |
| from . import Dim |
| |
| if isinstance(a, Dim) != isinstance(b, Dim): |
| return False |
| if isinstance(a, Dim): |
| return a is b |
| else: |
| return a == b |
| |
| |
| class isin: |
| def __contains__(self, item): |
| for x in self: |
| if seq(item, x): |
| return True |
| return False |
| |
| def index(self, item): |
| for i, x in enumerate(self): |
| if seq(item, x): |
| return i |
| raise ValueError |
| |
| |
| class llist(isin, list): |
| pass |
| |
| |
| class ltuple(isin, tuple): |
| pass |
| |
| |
| empty_dict = {} |
| |
| |
| @classmethod |
| def __torch_function__(self, orig, cls, args, kwargs=empty_dict): |
| from . import _Tensor, Tensor, TensorLike |
| from .delayed_mul_tensor import DelayedMulTensor |
| |
| if orig is torch.Tensor.__mul__: |
| lhs, rhs = args |
| if ( |
| isinstance(lhs, _Tensor) |
| and isinstance(rhs, _Tensor) |
| and lhs.ndim == 0 |
| and rhs.ndim == 0 |
| ): |
| return DelayedMulTensor(lhs, rhs) |
| all_dims = llist() |
| flat_args, unflatten = tree_flatten((args, kwargs)) |
| device_holding_tensor = None |
| for f in flat_args: |
| if isinstance(f, _Tensor): |
| if f._has_device: |
| device_holding_tensor = f._batchtensor |
| for d in f.dims: |
| if d not in all_dims: |
| all_dims.append(d) |
| |
| def unwrap(t): |
| if isinstance(t, _Tensor): |
| r = t._batchtensor |
| if device_holding_tensor is not None and not t._has_device: |
| r = r.to(device=device_holding_tensor.device) |
| return r |
| return t |
| |
| if orig in pointwise: |
| result_levels = llist() |
| arg_levels = llist() |
| to_expand = [] |
| for i, f in enumerate(flat_args): |
| if isinstance(f, TensorLike): |
| ptensor, levels, _ = _tensor_levels(f) |
| if ( |
| isinstance(f, _Tensor) |
| and not f._has_device |
| and device_holding_tensor is not None |
| ): |
| ptensor = ptensor.to(device=device_holding_tensor.device) |
| flat_args[i] = ptensor |
| for l in levels: |
| if l not in result_levels: |
| result_levels.append(l) |
| to_expand.append((i, levels)) |
| |
| for i, levels in to_expand: |
| flat_args[i] = _match_levels(flat_args[i], levels, result_levels) |
| args, kwargs = unflatten(flat_args) |
| result = orig(*args, **kwargs) |
| |
| def wrap(t): |
| if isinstance(t, TensorLike): |
| return Tensor.from_positional( |
| t, result_levels, device_holding_tensor is not None |
| ) |
| return t |
| |
| return tree_map(wrap, result) |
| else: |
| |
| def wrap(t): |
| if isinstance(t, TensorLike): |
| return Tensor.from_batched(t, device_holding_tensor is not None) |
| return t |
| |
| with _enable_layers(all_dims): |
| print(f"batch_tensor for {orig}") |
| args, kwargs = unflatten(unwrap(f) for f in flat_args) |
| result = orig(*args, **kwargs) |
| # print("END", orig) |
| return tree_map(wrap, result) |
| |
| |
| def positional(self, *dims): |
| from . import Dim, DimensionBindError, Tensor |
| |
| ptensor, levels = self._tensor, llist(self._levels) |
| flat_dims = llist() |
| view = [] |
| needs_view = False |
| ndim = self.ndim |
| for d in dims: |
| if isinstance(d, DimList): |
| flat_dims.extend(d) |
| view.extend(e.size for e in d) |
| elif isinstance(d, Dim): |
| flat_dims.append(d) |
| view.append(d.size) |
| elif isinstance(d, int): |
| d = _wrap_dim(d, ndim, False) |
| flat_dims.append(d) |
| view.append(ptensor.size(d)) |
| else: |
| flat_dims.extend(d) |
| view.append(prod(e.size for e in d)) |
| needs_view = True |
| |
| permute = list(range(len(levels))) |
| nflat = len(flat_dims) |
| for i, d in enumerate(flat_dims): |
| try: |
| idx = levels.index(d) |
| except ValueError as e: |
| raise DimensionBindError( |
| f"tensor of dimensions {self.dims} does not contain dim {d}" |
| ) from e |
| p = permute[idx] |
| del levels[idx] |
| del permute[idx] |
| levels.insert(i, 0) |
| permute.insert(i, p) |
| ptensor = ptensor.permute(*permute) |
| seen = 0 |
| for i in range(len(levels) - 1, -1, -1): |
| if isinstance(levels[i], int): |
| seen += 1 |
| levels[i] = -seen |
| result = Tensor.from_positional(ptensor, levels, self._has_device) |
| if needs_view: |
| result = result.reshape(*view, *result.size()[len(flat_dims) :]) |
| return result |
| |
| |
| def _contains_dim(input): |
| from . import Dim |
| |
| for i in input: |
| if isinstance(i, Dim): |
| return True |
| |
| |
| def expand(self, *sizes): |
| if not _contains_dim(sizes): |
| return self.__torch_function__(torch.Tensor.expand, None, (self, *sizes)) |
| dims = sizes |
| sizes = [d.size for d in dims] + [-1] * self.ndim |
| self = self.expand(*sizes) |
| return self[dims] |
| |
| |
| _not_present = object() |
| |
| |
| def _getarg(name, offset, args, kwargs, default): |
| if len(args) > offset: |
| return args[offset] |
| return kwargs.get(name, default) |
| |
| |
| def _patcharg(name, offset, args, kwargs, value): |
| if len(args) > offset: |
| args[offset] = value |
| else: |
| kwargs[name] = value |
| |
| |
| def _wrap( |
| orig, dim_offset=0, keepdim_offset=1, dim_name="dim", single_dim=False, reduce=True |
| ): |
| from . import Dim, Tensor, TensorLike |
| |
| def fn(self, *args, **kwargs): |
| dim = _getarg(dim_name, dim_offset, args, kwargs, _not_present) |
| if dim is _not_present or (single_dim and not isinstance(dim, Dim)): |
| with _enable_layers(self.dims): |
| print(f"dim fallback batch_tensor for {orig}") |
| return Tensor.from_batched( |
| orig(self._batchtensor, *args, **kwargs), self._has_device |
| ) |
| keepdim = ( |
| _getarg("keepdim", keepdim_offset, args, kwargs, False) if reduce else False |
| ) |
| t, levels = self._tensor, llist(self._levels) |
| dims = _dims(dim, self._batchtensor.ndim, keepdim, single_dim) |
| dim_indices = tuple(levels.index(d) for d in dims) |
| if reduce and not keepdim: |
| new_levels = [l for i, l in enumerate(levels) if i not in dim_indices] |
| else: |
| new_levels = levels |
| |
| if len(dim_indices) == 1: |
| dim_indices = dim_indices[ |
| 0 |
| ] # so that dims that really only take a single argument work... |
| args = list(args) |
| _patcharg(dim_name, dim_offset, args, kwargs, dim_indices) |
| |
| def wrap(t): |
| if isinstance(t, TensorLike): |
| return Tensor.from_positional(t, new_levels, self._has_device) |
| return t |
| |
| with _enable_layers(new_levels): |
| print(f"dim used batch_tensor for {orig}") |
| r = orig(t, *args, **kwargs) |
| return tree_map(wrap, r) |
| |
| return fn |
| |
| |
| def _def(name, *args, **kwargs): |
| from . import _Tensor |
| |
| orig = getattr(torch.Tensor, name) |
| setattr(_Tensor, name, _wrap(orig, *args, **kwargs)) |
| |
| |
| no_slice = slice(None) |
| |
| _orig_getitem = torch.Tensor.__getitem__ |
| |
| |
| class dim_tracker: |
| def __init__(self) -> None: |
| self.dims = llist() |
| self.count = [] |
| |
| def record(self, d): |
| if d not in self.dims: |
| self.dims.append(d) |
| self.count.append(1) |
| |
| def __getitem__(self, d): |
| return self.count[self.dims.index(d)] |
| |
| |
| def t__getitem__(self, input): |
| from . import _Tensor, Dim, DimensionBindError, DimList, Tensor, TensorLike |
| |
| # * bail to original example if we have a single non-Dim tensor, or a non-tensor |
| # * locate ... or an unbound tensor list, and determine its size, bind dim list |
| # (remember that None does not count to the total dim count) |
| # * bind simple dims and dim-packs to their sizes, count the number of uses of each dim, |
| # produce the re-view if needed |
| # * for each single-use dim index, replace with no_slice and mark that it will be added |
| # (keep track of whether we have to call super) |
| # * call super if needed |
| # * if we have dims to bind, bind them (it will help if we eliminated ... and None before) |
| # this handles bool indexing handling, as well as some other simple cases. |
| |
| is_simple = ( |
| not isinstance(input, Dim) |
| and not isinstance(input, (tuple, list)) |
| and |
| # WAR for functorch bug where zero time tensors in getitem are not handled correctly. |
| not (isinstance(input, TensorLike) and input.ndim == 0) |
| ) |
| |
| if is_simple: |
| if isinstance(self, _Tensor): |
| return _Tensor.__torch_function__(_orig_getitem, None, (self, input)) |
| else: |
| return _orig_getitem(self, input) |
| |
| # can further optimize this case |
| if not isinstance(input, tuple): |
| input = [input] |
| else: |
| input = list(input) |
| |
| dims_indexed = 0 |
| expanding_object = None |
| dimlists = [] |
| for i, s in enumerate(input): |
| if s is ... or isinstance(s, DimList) and not s.is_bound: |
| if expanding_object is not None: |
| msg = ( |
| "at most one ... or unbound dimension list can exist in indexing list but" |
| f" found 2 at offsets {i} and {expanding_object}" |
| ) |
| raise DimensionBindError(msg) |
| expanding_object = i |
| |
| if isinstance(s, DimList): |
| dims_indexed += len(s) if s.is_bound else 0 |
| dimlists.append(i) |
| elif s is not None and s is not ...: |
| dims_indexed += 1 |
| |
| ndim = self.ndim |
| if dims_indexed > ndim: |
| raise IndexError( |
| f"at least {dims_indexed} indices were supplied but the tensor only has {ndim} dimensions." |
| ) |
| if expanding_object is not None: |
| expanding_ndims = ndim - dims_indexed |
| obj = input[expanding_object] |
| if obj is ...: |
| input[expanding_object : expanding_object + 1] = [ |
| no_slice |
| ] * expanding_ndims |
| else: |
| obj.bind_len(expanding_ndims) |
| # flatten the dimslists into the indexing |
| for i in reversed(dimlists): |
| input[i : i + 1] = input[i] |
| dims_indexed = 0 |
| requires_view = False |
| size = self.size() |
| view_sizes = [] |
| dims_seen = dim_tracker() |
| |
| def add_dims(t): |
| if not isinstance(t, _Tensor): |
| return |
| for d in t.dims: |
| dims_seen.record(d) |
| |
| add_dims(self) |
| dim_packs = [] |
| for i, idx in enumerate(input): |
| if idx is None: |
| input[i] = no_slice |
| view_sizes.append(1) |
| requires_view = True |
| else: |
| sz = size[dims_indexed] |
| if isinstance(idx, Dim): |
| idx.size = sz |
| dims_seen.record(idx) |
| view_sizes.append(sz) |
| elif isinstance(idx, (tuple, list)) and idx and isinstance(idx[0], Dim): |
| for d in idx: |
| dims_seen.record(idx) |
| _bind_dims_to_size(sz, idx, f"offset {i}") |
| view_sizes.extend(d.size for d in idx) |
| requires_view = True |
| dim_packs.append(i) |
| else: |
| add_dims(idx) |
| view_sizes.append(sz) |
| dims_indexed += 1 |
| if requires_view: |
| self = self.view(*view_sizes) |
| for i in reversed(dim_packs): |
| input[i : i + 1] = input[i] |
| |
| # currenty: |
| # input is flat, containing either Dim, or Tensor, or something valid for standard indexing |
| # self may have first-class dims as well. |
| |
| # to index: |
| # drop the first class dims from self, they just become direct indices of their positions |
| |
| # figure out the dimensions of the indexing tensors: union of all the dims in the tensors in the index. |
| # these dimensions will appear and need to be bound at the first place tensor occures |
| |
| if isinstance(self, _Tensor): |
| ptensor_self, levels = self._tensor, list(self._levels) |
| # indices to ptensor rather than self which has first-class dimensions |
| input_it = iter(input) |
| flat_inputs = [next(input_it) if isinstance(l, int) else l for l in levels] |
| has_device = self._has_device |
| to_pad = 0 |
| else: |
| ptensor_self, flat_inputs = self, input |
| to_pad = ptensor_self.ndim - len(flat_inputs) |
| has_device = True |
| |
| result_levels = [] |
| index_levels = [] |
| tensor_insert_point = None |
| to_expand = {} |
| requires_getindex = False |
| for i, inp in enumerate(flat_inputs): |
| if isinstance(inp, Dim) and dims_seen[inp] == 1: |
| flat_inputs[i] = no_slice |
| result_levels.append(inp) |
| elif isinstance(inp, TensorLike): |
| requires_getindex = True |
| if tensor_insert_point is None: |
| tensor_insert_point = len(result_levels) |
| ptensor, levels, _ = _tensor_levels(inp) |
| to_expand[i] = levels |
| flat_inputs[i] = ptensor |
| for l in levels: |
| if l not in index_levels: |
| index_levels.append(l) |
| else: |
| requires_getindex = True |
| result_levels.append(0) |
| |
| if tensor_insert_point is not None: |
| result_levels[tensor_insert_point:tensor_insert_point] = index_levels |
| |
| for i, levels in to_expand.items(): |
| flat_inputs[i] = _match_levels(flat_inputs[i], levels, index_levels) |
| |
| if requires_getindex: |
| result = _orig_getitem(ptensor_self, flat_inputs) |
| else: |
| result = ptensor_self |
| |
| next_positional = -1 |
| if to_pad > 0: |
| result_levels.extend([0] * to_pad) |
| for i, r in enumerate(reversed(result_levels)): |
| if isinstance(r, int): |
| result_levels[-1 - i] = next_positional |
| next_positional -= 1 |
| |
| return Tensor.from_positional(result, result_levels, has_device) |
| |
| |
| # XXX - dim is optional and can be the outer-most dimension... |
| def stack(tensors, new_dim, dim=0, out=None): |
| if isinstance(dim, int): |
| return torch.stack(tensors, dim, out).index(dim, new_dim) |
| index = None |
| if out is not None: |
| out, index = _positional_no_permute(out, dim, expand_dim=True) |
| ptensors = [] |
| for t in tensors: |
| pt, pi = _positional_no_permute(t, dim, expand_dim=True) |
| if index is not None and pi != index: |
| pt = pt.move_dim(pi, index) |
| else: |
| index = pi |
| ptensors.append(pt) |
| pr = torch.stack(ptensors, index, out=out) |
| return pr.index((index, index + 1), (new_dim, dim)) |
| |
| |
| _orig_split = torch.Tensor.split |
| |
| |
| def split(self, split_size_or_sections, dim=0): |
| from . import _Tensor, Dim |
| |
| if isinstance(split_size_or_sections, int) or any( |
| isinstance(t, int) for t in split_size_or_sections |
| ): |
| if isinstance(dim, Dim): |
| raise ValueError( |
| "when dim is specified as a Dim object, split sizes must also be dimensions." |
| ) |
| return _orig_split(self, split_size_or_sections, dim=dim) |
| |
| if isinstance(dim, Dim): |
| assert isinstance(self, _Tensor), f"Tensor does not have dimension {dim}" |
| self, dim = _positional_no_permute(self, dim) |
| |
| size = self.size(dim) |
| total_bound_size = 0 |
| unbound = [] |
| sizes = [] |
| for i, d in enumerate(split_size_or_sections): |
| if d.is_bound: |
| sizes.append(d.size) |
| total_bound_size += d.size |
| else: |
| sizes.append(0) |
| unbound.append(i) |
| |
| if unbound: |
| assert ( |
| total_bound_size <= size |
| ), f"result dimensions are larger than original: {total_bound_size} vs {size} ({split_size_or_sections})" |
| remaining_size = size - total_bound_size |
| chunk_size = -(-remaining_size // len(unbound)) |
| for u in unbound: |
| sz = min(chunk_size, remaining_size) |
| split_size_or_sections[u].size = sz |
| sizes[u] = sz |
| remaining_size -= sz |
| else: |
| assert ( |
| total_bound_size == size |
| ), f"result dimensions do not match original: {total_bound_size} vs {size} ({split_size_or_sections})" |
| return tuple( |
| t.index(dim, d) |
| for d, t in zip(split_size_or_sections, _orig_split(self, sizes, dim=dim)) |
| ) |
