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)))