// TODO: How to type the "sub-functions"?
// TODO: improve this and make array types more specific
InTArray = TypeVar("InTArray", bound=Array)
InputLengths = TypeVar("InputLengths", bound=Array)
InTList = List[InTArray]
InTRagged = Tuple[InTArray, InputLengths]
InT = Union[InTArray, InTList, InTRagged]
OutTArray = TypeVar("OutTArray", bound=Array)
OutputLengths = TypeVar("OutputLengths", bound=Array)
OutTList = List[OutTArray]
After Change
// TODO: improve this and make array types more specific?
InT = Union[Array, List[Array], Ragged]
OutT = Union[Array, List[Array], Ragged]