Spaces:
Running
Running
| # 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 = {} | |
| 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): | |
| 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)) | |
| ) | |