avoid nesting then flattening

This commit is contained in:
kadarakos 2023-04-26 19:10:23 +00:00
parent 8d4129e177
commit 0c7c503a90

View File

@ -57,9 +57,9 @@ def _get_span_indices(ops, spans: Ragged, lengths: Ints1d) -> Ints1d:
for i, length in enumerate(lengths):
spans_i = spans[i].dataXd + offset
for j in range(spans_i.shape[0]):
indices.append(ops.xp.arange(spans_i[j, 0], spans_i[j, 1])) # type: ignore[call-overload, index]
indices += list(range(spans_i[j, 0], spans_i[j, 1])) # type: ignore[call-overload, index]
offset += length
return ops.flatten(indices, dtype="i", ndim_if_empty=1)
return ops.xp.asarray(indices, dtype="int")
def _ensure_cpu(spans: Ragged, lengths: Ints1d) -> Tuple[Ragged, Ints1d]: