680f4831e6c8f3102be8d99056b746b48ce34051,einops.py,,reduce,#,243

Before Change

    for composite_axis in composite_axes_left:
        found = {axis for axis in composite_axis if known_lengths[axis] is not None}
        not_found = {axis for axis in composite_axis if known_lengths[axis] is None}
        if len(not_found) == 0:
            // imitating that size of the first one was not computed
            not_found_axis = composite_axis[0]
            found.remove(not_found_axis)
            not_found.add(not_found_axis)

        assert len(not_found) == 1
        assert len(not_found) + len(found) == len(composite_axis)
        axis, = not_found
        computed_id = position_lookup[axis]

After Change

        lookup = dict(zip(list(known_lengths), range(len(known_lengths))))
        assert len(not_found) <= 1
        assert len(not_found) + len(found) == len(composite_axis)
        input_axes_known_unknown.append(([lookup[axis] for axis in found], [lookup[axis] for axis in not_found]))

    result_axes_grouping = [[position_lookup_after_reduction[axis] for axis in composite_axis]
                            for composite_axis in composite_axes_rght]

Instances