Repository: pfnet/pytorch-pfn-extras

Issue:
Module to validate a shape and dtype?
Modules like `tf.ensure_shape` can be useful in PyTorch too.
https://www.tensorflow.org/api_docs/python/tf/ensure_shape


Discussion:

Resolving patch:
diff --git a/example/mnist.py b/example/mnist.py
--- a/example/mnist.py
+++ b/example/mnist.py
@@ -25,6 +25,7 @@ def forward(self, x):
         x = x.flatten(start_dim=1)
         x = F.relu(self.fc1(x))
         x = self.fc2(x)
+        ppe.nn.ensure(x, shape=(None, 10))
         return F.log_softmax(x, dim=1)
 
 
diff --git a/pytorch_pfn_extras/nn/__init__.py b/pytorch_pfn_extras/nn/__init__.py
--- a/pytorch_pfn_extras/nn/__init__.py
+++ b/pytorch_pfn_extras/nn/__init__.py
@@ -1,3 +1,4 @@
+from pytorch_pfn_extras.nn.modules.ensure_shape import Ensure, ensure  # NOQA
 from pytorch_pfn_extras.nn.modules.lazy_linear import LazyLinear  # NOQA
 from pytorch_pfn_extras.nn.modules.lazy_conv import LazyConv1d  # NOQA
 from pytorch_pfn_extras.nn.modules.lazy_conv import LazyConv2d  # NOQA
diff --git a/pytorch_pfn_extras/nn/modules/ensure_shape.py b/pytorch_pfn_extras/nn/modules/ensure_shape.py
new file mode 100644
--- /dev/null
+++ b/pytorch_pfn_extras/nn/modules/ensure_shape.py
@@ -0,0 +1,115 @@
+from typing import Optional, Tuple
+
+import torch
+
+
+class Ensure(torch.nn.Module):
+    """Module to check the shape of a tensor.
+
+    Args:
+       shape: Tuple with the desired shape. If the input tensor shape
+           is not compatible, `ValueError` will be raised. If `None` is set
+           as a dimension value, that dimension will be ignored.
+       dtype: Checks if the `dtype` of the input thensor matches the
+           provided one.
+       broadcastable: Check if the shapes are compatible using broadcasting
+           rules.
+       can_cast: Check if the input tensor can be casted to the provided type.
+    """
+
+    def __init__(
+            self,
+            *,
+            shape: Optional[Tuple[Optional[int], ...]] = None,
+            dtype: Optional[torch.dtype] = None,
+            broadcastable: bool = False,
+            can_cast: bool = False,
+    ):
+        super().__init__()  # type: ignore[no-untyped-call]
+        if shape is None and dtype is None:
+            raise ValueError(
+                'shape, dtype or both arguments must be specified')
+        self._dtype = dtype
+        self._broadcastable = broadcastable
+        self._can_cast = can_cast
+        # Check if there are Nones in the shape and replace them by 1s
+        # so we can compare the shapes using broadcast semantics
+        c_shape: Optional[Tuple[int, ...]] = None
+        if shape is not None:
+            non_none_tuple = tuple(
+                [x if x is not None else 1
+                 for x in shape]
+            )
+            if None in shape:
+                self._broadcastable = True
+            c_shape = non_none_tuple
+
+        self._shape = None
+        if c_shape is not None:
+            # This is required for torch script
+            self._shape = torch.tensor(c_shape)
+
+    def _broadcast(self, shape_1: torch.Tensor, shape_2: torch.Tensor) -> bool:
+        # Torch broadcast_shapes raises an exception, we want a simple
+        # method that returs True/False so we can use torch script
+        l_1, l_2 = len(shape_1), len(shape_2)
+        if l_1 != l_2:
+            # we would like to do checks like the below one
+            # but torchscript converts everything to tensors
+            # shape_2 = (1,) * (l_1 - l_2) + shape_2
+            if l_1 < l_2:
+                shape_1 = torch.cat([torch.ones(l_2 - l_1), shape_1])
+            else:
+                shape_2 = torch.cat([torch.ones(l_1 - l_2), shape_2])
+        # Look for Elements that are different in the tensors
+        for a, b in zip(shape_1, shape_2):
+            if a != b and a != 1 and b != 1:
+                return False
+        return True
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        # To make it compatible with torchscript since torch.Size does not work
+        t_shape = torch.tensor(input.shape)
+        if self._shape is not None and list(t_shape) != list(self._shape):
+            if self._broadcastable:
+                if not self._broadcast(t_shape, self._shape):
+                    raise ValueError(
+                        f'Shapes {self._shape} and {input.shape} are non'
+                        ' broadcastable')
+            else:
+                raise ValueError(
+                    f'Expected {self._shape}, input shape is {input.shape}')
+        if self._dtype is not None and input.dtype != self._dtype:
+            if self._can_cast:
+                if not torch.can_cast(input.dtype, self._dtype):
+                    raise ValueError(
+                        f'Input dtype {input.dtype} can\'t be casted to'
+                        f' {self._dtype}')
+            else:
+                raise ValueError(
+                    f'Expected {self._dtype}, input dtype is {input.dtype}')
+        return input
+
+
+def ensure(
+        tensor: torch.Tensor,
+        shape: Optional[Tuple[Optional[int], ...]] = None,
+        dtype: Optional[torch.dtype] = None,
+        broadcastable: bool = False,
+        can_cast: bool = False
+) -> None:
+    """Checks the shape and type of a tensor.
+
+    Args:
+       shape: Tuple with the desired shape. If the input tensor shape
+           is not compatible, `ValueError` will be raised. If `None` is set
+           as a dimension value, that dimension will be ignored.
+       dtype: Checks if the `dtype` of the input thensor matches the
+           provided one.
+       broadcastable: Check if the shapes are compatible using broadcasting
+           rules.
+       can_cast: Check if the input tensor can be casted to the provided type.
+    """
+    Ensure(
+        shape=shape, dtype=dtype,
+        broadcastable=broadcastable, can_cast=can_cast)(tensor)


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_ensure_shape.py b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_ensure_shape.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_ensure_shape.py
@@ -0,0 +1,164 @@
+import pytest
+import torch
+
+from pytorch_pfn_extras.nn import Ensure, ensure
+
+
+class TestEnsure:
+    def test_wrong_initialization(self):
+        with pytest.raises(ValueError, match='both arguments'):
+            Ensure(shape=None, dtype=None)
+
+    @pytest.mark.parametrize(
+        'shape', [(), (1,), (1, 1), (2,), (2, 4), (2, 3, 4)]
+    )
+    def test_valid_shape(self, shape):
+        tensor = torch.zeros(shape)
+        module = Ensure(shape=shape)
+        module(tensor)
+        # Use the function version
+        ensure(tensor, shape)
+
+    @pytest.mark.parametrize(
+        'shape', [(), (1,), (1, 1), (2,), (2, 4), (2, 3, 4)]
+    )
+    def test_invalid_shape(self, shape):
+        tensor = torch.zeros((1, 2, 3))
+        module = Ensure(shape=shape)
+        with pytest.raises(ValueError, match='input shape is'):
+            module(tensor)
+        with pytest.raises(ValueError, match='input shape is'):
+            ensure(tensor, shape)
+
+    @pytest.mark.parametrize('shape_t, shape_c', [
+        ((), (2,)),
+        ((3,), ()),
+        ((1,), (2,)),
+        ((1, 1), (2, 1)),
+        ((1, 1), (2,)),
+        ((2, 1), (2, 2)),
+        ((2, 4), (4,)),
+        ((2, 3, 4), (1, 4)),
+    ])
+    def test_broadcastable_shape(self, shape_t, shape_c):
+        tensor = torch.zeros(shape_t)
+        module = Ensure(shape=shape_c, broadcastable=True)
+        module(tensor)
+
+    @pytest.mark.parametrize('shape_t, shape_c', [
+        ((3,), (2,)),
+        ((2, 3), (2, 2)),
+        ((2, 4), (3, 4)),
+        ((2, 4), (3,)),
+        ((2, 3, 4), (2, 2, 1)),
+    ])
+    def test_nonbroadcastable_shape(self, shape_t, shape_c):
+        tensor = torch.zeros(shape_t)
+        module = Ensure(shape=shape_c, broadcastable=True)
+        with pytest.raises(ValueError, match='non broadcastable'):
+            module(tensor)
+
+    @pytest.mark.parametrize('shape_t, shape_c', [
+        ((2,), (None,)),
+        ((2, 2), (2, None)),
+        ((2, 1), (None, 2)),
+        ((1, 4), (None, 4)),
+        ((2, 3, 4), (2, None, 4)),
+    ])
+    def test_unknown_shape(self, shape_t, shape_c):
+        tensor = torch.zeros(shape_t)
+        module = Ensure(shape=shape_c)
+        module(tensor)
+
+    @pytest.mark.parametrize('shape_t, shape_c', [
+        ((3, 2), (2, None)),
+        ((1, 4), (None, 2)),
+        ((2, 3, 4), (3, None, 4)),
+    ])
+    def test_invalid_unknown_shape(self, shape_t, shape_c):
+        tensor = torch.zeros(shape_t)
+        module = Ensure(shape=shape_c)
+        with pytest.raises(ValueError, match='non broadcastable'):
+            module(tensor)
+
+    @pytest.mark.parametrize(
+        'dtype', [torch.int32, torch.float32, torch.complex64]
+    )
+    def test_valid_dtypes(self, dtype):
+        tensor = torch.zeros(1, dtype=dtype)
+        module = Ensure(shape=None, dtype=dtype)
+        module(tensor)
+        ensure(tensor, None, dtype)
+
+    @pytest.mark.parametrize('dtype_t, dtype_c', [
+        (torch.int32, torch.int16),
+        (torch.int32, torch.float32),
+        (torch.float32, torch.float64),
+        (torch.float32, torch.complex64),
+    ])
+    def test_invalid_dtypes(self, dtype_t, dtype_c):
+        tensor = torch.zeros(1, dtype=dtype_t)
+        module = Ensure(shape=None, dtype=dtype_c)
+        with pytest.raises(ValueError, match='input dtype'):
+            module(tensor)
+
+    @pytest.mark.parametrize('dtype_t, dtype_c', [
+        (torch.int32, torch.float32),
+        (torch.int32, torch.complex128),
+        (torch.int8, torch.float16),
+    ])
+    def test_dtypes_with_cast(self, dtype_t, dtype_c):
+        tensor = torch.zeros(1, dtype=dtype_t)
+        module = Ensure(shape=None, dtype=dtype_c, can_cast=True)
+        module(tensor)
+
+    @pytest.mark.parametrize('dtype_t, dtype_c', [
+        (torch.complex64, torch.int32),
+        (torch.float32, torch.int32),
+    ])
+    def test_invalid_dtypes_with_cast(self, dtype_t, dtype_c):
+        tensor = torch.zeros(1, dtype=dtype_t)
+        module = Ensure(shape=None, dtype=dtype_c, can_cast=True)
+        with pytest.raises(ValueError, match='be casted to'):
+            module(tensor)
+
+    def test_valid_shape_and_dtype(self):
+        shape = (10, 5)
+        dtype = torch.float32
+        tensor = torch.zeros(shape, dtype=dtype)
+        module = Ensure(shape=shape, dtype=dtype)
+        module(tensor)
+        ensure(tensor, shape, dtype)
+
+    # Too many warnings to list them all
+    @pytest.mark.filterwarnings('ignore')
+    def test_jit_module(self):
+        shape = (10, 5)
+        dtype = torch.float32
+        tensor = torch.zeros(shape, dtype=dtype)
+        module = Ensure(shape=shape, dtype=dtype)
+        jit_module = torch.jit.trace(module, (tensor,))
+        jit_module(tensor)
+
+        # An invalid tensor also passes, because checks are disabled
+        shape = (5, 5)
+        tensor = torch.zeros(shape, dtype=dtype)
+        jit_module(tensor)
+
+        # Tracing with a different shape fails during trace process
+        with pytest.raises(ValueError, match='input shape is'):
+            jit_module = torch.jit.trace(module, (tensor,))
+
+    def test_torchscript_module(self):
+        shape = (10, 5)
+        dtype = torch.float32
+        tensor = torch.zeros(shape, dtype=dtype)
+        module = Ensure(shape=shape, dtype=dtype)
+        jit_module = torch.jit.script(module)
+        jit_module(tensor)
+
+        shape = (5, 5)
+        tensor = torch.zeros(shape, dtype=dtype)
+        # torchscript changes the exception type
+        with pytest.raises(torch.jit.Error, match='input shape is'):
+            jit_module(tensor)