cleanup

#116

Author: albertz

nn/conv.py

@@ -69,12 +69,10 @@ def _lazy_init(self, in_dim: nn.Dim):
 
   def _call_nd1(self, source: nn.Tensor, *,
                 in_dim: Optional[nn.Dim] = None,
-                in_spatial_dim: nn.Dim,
-                out_spatial_dim: Optional[nn.Dim] = None) -> Tuple[nn.Tensor, nn.Dim]:
+                in_spatial_dim: nn.Dim) -> Tuple[nn.Tensor, nn.Dim]:
     assert self.nd == 1
     out, (out_spatial_dim,) = self.__class__.__base__.__call__(
-      self, source, in_dim=in_dim, in_spatial_dims=[in_spatial_dim],
-      out_spatial_dims=[out_spatial_dim] if out_spatial_dim else None)
+      self, source, in_dim=in_dim, in_spatial_dims=[in_spatial_dim])
     return out, out_spatial_dim
 
 
@@ -119,21 +117,19 @@ def __init__(self,
   @nn.scoped
   def __call__(self, source: nn.Tensor, *,
                in_dim: Optional[nn.Dim] = None,
-               in_spatial_dims: Sequence[nn.Dim],
-               out_spatial_dims: Optional[Sequence[nn.Dim]] = None
+               in_spatial_dims: Sequence[nn.Dim]
                ) -> Tuple[nn.Tensor, Sequence[nn.Dim]]:
     source = nn.check_in_feature_dim_lazy_init(source, in_dim, self.in_dim, self._lazy_init)
     for in_spatial_dim in in_spatial_dims:
       if in_spatial_dim not in source.shape:
         raise ValueError(f"{self}: source {source} does not have spatial dim {in_spatial_dim}")
-    if not out_spatial_dims:
-      out_spatial_dims = _default_out_spatial_dims(
-        description_prefix=nn.NameCtx.current_ctx().layer_abs_name_scope,
-        in_spatial_dims=in_spatial_dims,
-        filter_size=[d.dimension for d in self.filter_size],
-        strides=1 if not self.strides else self.strides,
-        dilation_rate=1 if not self.dilation_rate else self.dilation_rate,
-        padding=self.padding)
+    out_spatial_dims = _default_out_spatial_dims(
+      description_prefix=nn.NameCtx.current_ctx().layer_abs_name_scope,
+      in_spatial_dims=in_spatial_dims,
+      filter_size=[d.dimension for d in self.filter_size],
+      strides=1 if not self.strides else self.strides,
+      dilation_rate=1 if not self.dilation_rate else self.dilation_rate,
+      padding=self.padding)
     layer_dict = {
       "class": "conv", "from": source,
       "in_dim": self.in_dim, "in_spatial_dims": in_spatial_dims,
@@ -242,18 +238,16 @@ def __init__(self,
   @nn.scoped
   def __call__(self, source: nn.Tensor, *,
                in_dim: Optional[nn.Dim] = None,
-               in_spatial_dims: Sequence[nn.Dim],
-               out_spatial_dims: Optional[Sequence[nn.Dim]] = None
+               in_spatial_dims: Sequence[nn.Dim]
                ) -> Tuple[nn.Tensor, Sequence[nn.Dim]]:
     source = nn.check_in_feature_dim_lazy_init(source, in_dim, self.in_dim, self._lazy_init)
-    if not out_spatial_dims:
-      out_spatial_dims = [
-        nn.SpatialDim(f"{nn.NameCtx.current_ctx().layer_abs_name_scope}:out-spatial-dim{i}")
-        for i, s in enumerate(self.filter_size)]
-      for i in range(len(self.filter_size)):
-        s = self.filter_size[i].dimension if not self.strides else self.strides[i]
-        if self.filter_size[i].dimension == s == 1 or (s == 1 and self.padding.lower() == "same"):
-          out_spatial_dims[i] = in_spatial_dims[i]
+    out_spatial_dims = [
+      nn.SpatialDim(f"{nn.NameCtx.current_ctx().layer_abs_name_scope}:out-spatial-dim{i}")
+      for i, s in enumerate(self.filter_size)]
+    for i in range(len(self.filter_size)):
+      s = self.filter_size[i].dimension if not self.strides else self.strides[i]
+      if self.filter_size[i].dimension == s == 1 or (s == 1 and self.padding.lower() == "same"):
+        out_spatial_dims[i] = in_spatial_dims[i]
     layer_dict = {
       "class": "transposed_conv", "from": source,
       "in_dim": self.in_dim, "in_spatial_dims": in_spatial_dims,