Avoid h2d - d2h roundtrip when using unflatten

thinc/layers/concatenate.py

@@ -1,5 +1,7 @@
 from typing import Any, List, Tuple, Callable, Optional
 from typing import TypeVar, cast, Dict, Union, Sequence
+
+from ..backends import NumpyOps
 from ..model import Model
 from ..config import registry
 from ..types import Array2d, Ragged
@@ -8,6 +10,9 @@
 from ..types import XY_XY_OutT
 
 
+NUMPY_OPS = NumpyOps()
+
+
 InT = TypeVar("InT", bound=Any)
 OutT = TypeVar("OutT", bound=Union[Array2d, Sequence[Array2d], Ragged])
 
@@ -120,7 +125,7 @@ def backprop(d_output: Sequence[Array2d]) -> InT:
             start += width
         return dX
 
-    lengths = model.ops.asarray1i([len(x) for x in X])
+    lengths = NUMPY_OPS.asarray1i([len(x) for x in X])
     Ys = [model.ops.xp.concatenate(Y, axis=0) for Y in Ys]
     widths = [Y.shape[1] for Y in Ys]
     out_array = model.ops.xp.hstack(Ys)

thinc/layers/list2array.py

@@ -1,10 +1,14 @@
-from typing import Tuple, Callable, TypeVar, List, Union, cast
+from typing import Tuple, Callable, TypeVar, List
 
+from ..backends import NumpyOps
 from ..model import Model
 from ..config import registry
 from ..types import Array2d
 
 
+NUMPY_OPS = NumpyOps()
+
+
 OutT = TypeVar("OutT", bound=Array2d)
 InT = List[OutT]
 
@@ -19,7 +23,7 @@ def list2array() -> Model[InT, OutT]:
 
 
 def forward(model: Model[InT, OutT], Xs: InT, is_train: bool) -> Tuple[OutT, Callable]:
-    lengths = model.ops.asarray1i([len(x) for x in Xs])
+    lengths = NUMPY_OPS.asarray1i([len(x) for x in Xs])
 
     def backprop(dY: OutT) -> InT:
         return model.ops.unflatten(dY, lengths)

thinc/layers/with_array.py

@@ -1,9 +1,14 @@
 from typing import Tuple, Callable, Optional, TypeVar, Union, cast
 
+from ..backends import NumpyOps
 from ..model import Model
 from ..config import registry
 from ..types import Padded, Ragged, ArrayXd, Array3d, ListXd
 
+
+NUMPY_OPS = NumpyOps()
+
+
 ArrayTXd = TypeVar("ArrayTXd", bound=ArrayXd)
 SeqT = TypeVar("SeqT", bound=Union[Padded, Ragged, ListXd, ArrayXd])
 
@@ -68,7 +73,7 @@ def _list_forward(
 ) -> Tuple[ListXd, Callable]:
     layer: Model[ArrayXd, ArrayXd] = model.layers[0]
     pad = model.attrs["pad"]
-    lengths = layer.ops.asarray1i([len(seq) for seq in Xs])
+    lengths = NUMPY_OPS.asarray1i([len(seq) for seq in Xs])
     Xf = layer.ops.flatten(Xs, pad=pad)
     Yf, get_dXf = layer(Xf, is_train)
 

thinc/layers/with_array2d.py

@@ -1,10 +1,14 @@
 from typing import Tuple, Callable, Optional, TypeVar, cast, List, Union
 
+from ..backends import NumpyOps
 from ..model import Model
 from ..config import registry
 from ..types import Array2d, Floats2d, List2d, Padded, Ragged
 
 
+NUMPY_OPS = NumpyOps()
+
+
 ValT = TypeVar("ValT", bound=Array2d)
 SeqT = TypeVar("SeqT", bound=Union[Padded, Ragged, List2d, Array2d])
 
@@ -71,7 +75,7 @@ def _list_forward(
 ) -> Tuple[List2d, Callable]:
     layer: Model[Array2d, Array2d] = model.layers[0]
     pad = model.attrs["pad"]
-    lengths = layer.ops.asarray1i([len(seq) for seq in Xs])
+    lengths = NUMPY_OPS.asarray1i([len(seq) for seq in Xs])
     Xf = layer.ops.flatten(Xs, pad=pad)
     Yf, get_dXf = layer(Xf, is_train)
 

thinc/layers/with_ragged.py

@@ -1,9 +1,14 @@
 from typing import Tuple, Callable, Optional, TypeVar, cast, List, Union
 
+from ..backends import NumpyOps
 from ..types import Padded, Ragged, Array2d, ListXd, List2d, Ints1d
 from ..model import Model
 from ..config import registry
 
+
+NUMPY_OPS = NumpyOps()
+
+
 RaggedData = Tuple[Array2d, Ints1d]
 SeqT = TypeVar("SeqT", bound=Union[Padded, Ragged, ListXd, RaggedData])
 
@@ -89,9 +94,12 @@ def _padded_forward(
     # sooner.
     Xs = padded2list(Xp)
     # Bit annoying here: padded is in a different order, so we need to make new
-    # lengths.
-    lengths = layer.ops.asarray1i([len(x) for x in Xs])
-    Yr, get_dXr = layer(Ragged(flatten(Xs), lengths), is_train)
+    # lengths. The lengths are unconditionally allocated in CPU memory, because
+    # otherwire unflatten would move GPU allocations to the CPU again. For the
+    # ragged arrays we let the layer's ops determine how lengths should be
+    # stored to ensure that the array and lengths use the same type of memory.
+    lengths = NUMPY_OPS.asarray1i([len(x) for x in Xs])
+    Yr, get_dXr = layer(Ragged(flatten(Xs), layer.ops.asarray1i(lengths)), is_train)
 
     def backprop(dYp: Padded):
         flattened = flatten(padded2list(dYp))
@@ -111,8 +119,8 @@ def _list_forward(
     flatten = layer.ops.flatten
     unflatten = layer.ops.unflatten
 
-    lengths = layer.ops.asarray1i([len(x) for x in Xs])
-    Yr, get_dXr = layer(Ragged(flatten(Xs), lengths), is_train)
+    lengths = [len(x) for x in Xs]
+    Yr, get_dXr = layer(Ragged(flatten(Xs), layer.ops.asarray1i(lengths)), is_train)
 
     def backprop(dYs):
         flattened = flatten(dYs)