mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
eace053398
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/38211 Just because the annotations are inline doesn't mean the files type check; most of the newly annotated files have type errors and I added exclusions for them in mypy.ini. The payoff of moving all of these modules inline is I can delete the relevant code generation logic for the pyi files (which was added ignore annotations that weren't actually relevant anymore.) For the most part the translation was completely mechanical, but there were two hairy issues. First, I needed to work around a Python 3.6 and earlier bug where Generic has a nontrivial metaclass. This fix is in torch/jit/__init__.py. Second, module.py, we need to apply the same fix for avoiding contravariance checks that the pyi file used to have; this is done by declaring forward as a variable (rather than a function), which appears to be sufficient enough to get mypy to not contravariantly check input arguments. Because we aren't actually typechecking these modules in most cases, it is inevitable that some of these type annotations are wrong. I slavishly copied the old annotations from the pyi files unless there was an obvious correction I could make. These annotations will probably need fixing up later. Signed-off-by: Edward Z. Yang <ezyang@fb.com> Test Plan: Imported from OSS Differential Revision: D21497397 Pulled By: ezyang fbshipit-source-id: 2b08bacc152c48f074e7edc4ee5dce1b77d83702
34 lines
921 B
Python
34 lines
921 B
Python
from .module import Module
|
|
|
|
from torch import Tensor
|
|
|
|
class Flatten(Module):
|
|
r"""
|
|
Flattens a contiguous range of dims into a tensor. For use with :class:`~nn.Sequential`.
|
|
Args:
|
|
start_dim: first dim to flatten (default = 1).
|
|
end_dim: last dim to flatten (default = -1).
|
|
|
|
Shape:
|
|
- Input: :math:`(N, *dims)`
|
|
- Output: :math:`(N, \prod *dims)` (for the default case).
|
|
|
|
|
|
Examples::
|
|
>>> m = nn.Sequential(
|
|
>>> nn.Conv2d(1, 32, 5, 1, 1),
|
|
>>> nn.Flatten()
|
|
>>> )
|
|
"""
|
|
__constants__ = ['start_dim', 'end_dim']
|
|
start_dim: int
|
|
end_dim: int
|
|
|
|
def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
|
|
super(Flatten, self).__init__()
|
|
self.start_dim = start_dim
|
|
self.end_dim = end_dim
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
return input.flatten(self.start_dim, self.end_dim)
|