alpaca/alpacka/data/ops.py
def nested_zip(xs):
    """Zips a list of pytrees.

    Inverse of nested_unzip.

    Example:
        nested_zip([((1, 2), 3), ((4, 5), 6)]) == (([1, 4], [2, 5]), [3, 6])
    """
    assert not _is_leaf(xs)
    assert xs
    if not _is_leaf(xs[0]):
        # Assert that the first levels of the zipped trees are the same.
        for x in xs:
            assert type(x) is type(xs[0]), (
                'Cannot zip pytrees of different types: '
                '{} and {}.'.format(type(x), type(xs[0]))
            )

    if _is_namedtuple_instance(xs[0]):
        return type(xs[0])(*nested_zip([tuple(x) for x in xs]))
    elif isinstance(xs[0], (list, tuple)):
        for x in xs:
            assert len(x) == len(xs[0]), (
                'Cannot zip sequences of different lengths: '
                '{} and {}'.format(len(x), len(xs[0]))
            )

        return type(xs[0])(
            nested_zip([x[i] for x in xs]) for i in range(len(xs[0]))
        )
    elif isinstance(xs[0], dict):
        return {k: nested_zip([x[k] for x in xs]) for k in xs[0].keys()}
    else:
        return xs
Zips a list of pytrees.

    Inverse of nested_unzip.

    Example:
        nested_zip([((1, 2), 3), ((4, 5), 6)]) == (([1, 4], [2, 5]), [3, 6])
    
Zips a list of pytrees.
Zips a list of pytrees .
def nested_zip(xs):
    assert not _is_leaf(xs)
    assert xs
    if not _is_leaf(xs[0]):
        for x in xs:
            assert type(x) is type(xs[0]), (
                'Cannot zip pytrees of different types: '
                '{} and {}.'.format(type(x), type(xs[0]))
            )
    if _is_namedtuple_instance(xs[0]):
        return type(xs[0])(*nested_zip([tuple(x) for x in xs]))
    elif isinstance(xs[0], (list, tuple)):
        for x in xs:
            assert len(x) == len(xs[0]), (
                'Cannot zip sequences of different lengths: '
                '{} and {}'.format(len(x), len(xs[0]))
            )
        return type(xs[0])(
            nested_zip([x[i] for x in xs]) for i in range(len(xs[0]))
        )
    elif isinstance(xs[0], dict):
        return {k: nested_zip([x[k] for x in xs]) for k in xs[0].keys()}
    else:
        return xs
def nested_zip ( xs ) : assert not _is_leaf ( xs ) assert xs if not _is_leaf ( xs [ 0 ] ) : for x in xs : assert type ( x ) is type ( xs [ 0 ] ) , ( 'Cannot zip pytrees of different types: ' '{} and {}.' . format ( type ( x ) , type ( xs [ 0 ] ) ) ) if _is_namedtuple_instance ( xs [ 0 ] ) : return type ( xs [ 0 ] ) ( * nested_zip ( [ tuple ( x ) for x in xs ] ) ) elif isinstance ( xs [ 0 ] , ( list , tuple ) ) : for x in xs : assert len ( x ) == len ( xs [ 0 ] ) , ( 'Cannot zip sequences of different lengths: ' '{} and {}' . format ( len ( x ) , len ( xs [ 0 ] ) ) ) return type ( xs [ 0 ] ) ( nested_zip ( [ x [ i ] for x in xs ] ) for i in range ( len ( xs [ 0 ] ) ) ) elif isinstance ( xs [ 0 ] , dict ) : return { k : nested_zip ( [ x [ k ] for x in xs ] ) for k in xs [ 0 ] . keys ( ) } else : return xs
Zips a list of pytrees.
Zips a list of pytrees .
"""Zips a list of pytrees.

    Inverse of nested_unzip.

    Example:
        nested_zip([((1, 2), 3), ((4, 5), 6)]) == (([1, 4], [2, 5]), [3, 6])
    """ # Assert that the first levels of the zipped trees are the same.
xs
examples
None