def test_nested_stack(self):
inp = [
{
'a': ([0, 1], 2),
'b': _TestNamedtuple(3)
},
{
'a': ([1, 2], 3),
'b': _TestNamedtuple(4)
},
]
out = {'a': ([[0, 1], [1, 2]], [2, 3]), 'b': _TestNamedtuple([3, 4])}
onp.testing.assert_equal(fastmath.nested_stack(inp), out)
def stack(x):
if not x:
return None
return fastmath.nested_stack(x)
def batch_stream(stream, batch_size):
while True:
yield trax_math.nested_stack(list(itertools.islice(stream,
batch_size)))
