Можешь сделать иначе. Придумай элемент-терминатор для плоского массива. Например ты знаешь, что NaN или Inf у тебя никогда не будут попадаться, тогда 2д массив вида:
arr = [
[1.0, 2.0, 3.0],
[1.0],
[]
]
Можно представить в плоском виде как:
arrFlattened = [
1.0, 2.0, 3.0, NaN,
1.0, NaN,
NaN,
NaN
]