Merge branch 'master' of https://github.com/explosion/thinc

Author: honnibal

thinc/layers/add.py

@@ -1,12 +1,12 @@
-from typing import Tuple, Callable, Optional, TypeVar, Dict
+from typing import Any, Tuple, Callable, Optional, TypeVar, Dict
 
 from ..model import Model
 from ..config import registry
 from ..types import ArrayXd, XY_XY_OutT
 from ..util import get_width
 
 
-InT = TypeVar("InT", bound=ArrayXd)
+InT = TypeVar("InT", bound=Any)
 OutT = TypeVar("OutT", bound=ArrayXd)
 
 
@@ -30,7 +30,7 @@ def add(
     return Model("add", forward, init=init, dims=dims, layers=layers)
 
 
-def forward(model: Model[InT, InT], X: InT, is_train: bool) -> Tuple[InT, Callable]:
+def forward(model: Model[InT, OutT], X: InT, is_train: bool) -> Tuple[OutT, Callable]:
     if not model.layers:
         return X, lambda dY: dY
     Y, first_callback = model.layers[0](X, is_train=is_train)
@@ -40,7 +40,7 @@ def forward(model: Model[InT, InT], X: InT, is_train: bool) -> Tuple[InT, Callab
         Y += layer_Y
         callbacks.append(layer_callback)
 
-    def backprop(dY: InT) -> InT:
+    def backprop(dY: InT) -> OutT:
         dX = first_callback(dY)
         for callback in callbacks:
             dX += callback(dY)
@@ -50,14 +50,20 @@ def backprop(dY: InT) -> InT:
 
 
 def init(
-    model: Model[InT, InT], X: Optional[InT] = None, Y: Optional[InT] = None
-) -> Model[InT, InT]:
+    model: Model[InT, OutT], X: Optional[InT] = None, Y: Optional[OutT] = None
+) -> Model[InT, OutT]:
     if X is not None:
         if model.has_dim("nI") is not False:
             model.set_dim("nI", get_width(X))
         for layer in model.layers:
             if layer.has_dim("nI") is not False:
                 layer.set_dim("nI", get_width(X))
+    if Y is not None:
+        if model.has_dim("nO") is not False:
+            model.set_dim("nO", get_width(Y))
+        for layer in model.layers:
+            if layer.has_dim("nO") is not False:
+                layer.set_dim("nO", get_width(Y))
     for layer in model.layers:
         layer.initialize(X=X, Y=Y)
     model.set_dim("nO", model.layers[0].get_dim("nO"))