from typing import Tuple, List, Callable, TypeVar

from ..model import Model
from ..types import Array


InputValue = TypeVar("InputValue", bound=Array)
InputType = List[InputValue]
OutputValue = TypeVar("OutputValue", bound=Array)
OutputLengths = TypeVar("OutputLengths", bound=Array)
OutputType = Tuple[OutputValue, OutputLengths]


def flatten_add_lengths() -> Model:
    """Transform sequences to ragged arrays if necessary. If sequences are
    already ragged, do nothing. A ragged array is a tuple (data, lengths),
    where data is the concatenated data.
    """
    return Model("flatten_add_lengths", forward)


def forward(
    model: Model, seqs: InputType, is_train: bool
) -> Tuple[OutputType, Callable]:
    lengths = model.ops.asarray([len(seq) for seq in seqs], dtype="i")

    def backprop(dY: OutputType) -> InputType:
        return model.ops.unflatten(dY, lengths)

    return (model.ops.flatten(seqs), lengths), backprop

@@ -70,16 +70,42 @@
 t Model%0A
+from ..data import Ragged%0A
 from ..t
@@ -188,81 +188,59 @@
 e = 
-List%5BInputValue%5D%0AOutputValue = TypeVar(%22OutputValue%22, bound=Array
+TypeVar(%22InputType%22, bound=List%5BInputValue%5D
 )%0AOutput
 Leng
@@ -235,23 +235,20 @@
 )%0AOutput
-Lengths
+Type
  = TypeV
@@ -257,23 +257,20 @@
 (%22Output
-Lengths
+Type
 %22, bound
@@ -274,87 +274,33 @@
 und=
-Array)%0AOutputType = Tuple%5BOutputValue, OutputLengths%5D%0A%0A%0Adef flatten_add_lengths
+Ragged)%0A%0A%0Adef list2ragged
 () -
@@ -529,27 +529,19 @@
 el(%22
-flatten_add_lengths
+list2ragged
 %22, f
@@ -581,19 +581,17 @@
  Model, 
-seq
+X
 s: Input
@@ -648,79 +648,8 @@
 le%5D:
-%0A    lengths = model.ops.asarray(%5Blen(seq) for seq in seqs%5D, dtype=%22i%22)
 %0A%0A  
@@ -665,16 +665,17 @@
 kprop(dY
+r
 : Output
@@ -731,18 +731,28 @@
 atten(dY
-, 
+r.data, dYr.
 lengths)
@@ -753,16 +753,81 @@
 ngths)%0A%0A
+    lengths = model.ops.asarray(%5Blen(x) for x in Xs%5D, dtype=%22i%22)%0A
     retu
@@ -829,16 +829,22 @@
  return 
+Ragged
 (model.o
@@ -854,19 +854,17 @@
 flatten(
-seq
+X
 s), leng

Update list2ragged