Repository: numpy/numpy
File: numpy/_core/einsumfunc.py (Python)

Code before:
        if sizes.setdefault(ix, d) != d:
            # set or check size
            raise ValueError(
                f"Index {ix} has mismatched sizes {sizes[ix]} and {d}."
            )
        right[ix] = True

    # now we classify the unique size > 1 indices only
    bat_inds = []  # appears on A, B, O
    con_inds = []  # appears on A, B, .
    a_keep = []  # appears on A, ., O
    b_keep = []  # appears on ., B, O
    # other indices (appearing on A or B only) will
    # be summed or traced out prior to the matmul
    for ix in left:
        if right.pop(ix, False):
            if ix in out:
                bat_inds.append(ix)
            else:
                con_inds.append(ix)
        elif ix in out:
            a_keep.append(ix)
    # now only indices unique to right remain
    for ix in right:
        if ix in out:
            b_keep.append(ix)

    if not con_inds:
        # contraction is pure multiplication, prepare inputs differently
        return _parse_eq_to_pure_multiplication(
            a_term, shape_a, b_term, shape_b, out
        )

    # only need the size one indices that appear in the output
    singletons = [ix for ix in out if ix in singletons]

    # take diagonal, remove any trivial axes and transpose left
    desired_a = "".join((*bat_inds, *a_keep, *con_inds))
    if a_term != desired_a:
        eq_a = f"{a_term}->{desired_a}"
    else:
        eq_a = None

    # take diagonal, remove any trivial axes and transpose right
    desired_b = "".join((*bat_inds, *con_inds, *b_keep))
    if b_term != desired_b:
        eq_b = f"{b_term}->{desired_b}"
    else:
        eq_b = None

    # then we want to reshape

Review comment (performance):
this could be simplified a bit now:

```suggestion
    b_keep.extend(set(out) & right.keys())
```

Or maybe it's an idea to define something like an `outset = set(out)` earlier on, so that all other `_ in out` O(n) calls can be replaced with `_ in outset` O(1) calls? Although I wonder if it'll actually be faster in practice :thinking: 

(feel free to ignore if you don't like this btw)

Code after:
        if sizes.setdefault(ix, d) != d:
            # set and check size
            raise ValueError(
                f"Index {ix} has mismatched sizes {sizes[ix]} and {d}."
            )
        right[ix] = True

    # now we classify the unique size > 1 indices only
    bat_inds = []  # appears on A, B, O
    con_inds = []  # appears on A, B, .
    a_keep = []  # appears on A, ., O
    b_keep = []  # appears on ., B, O
    # other indices (appearing on A or B only) will
    # be summed or traced out prior to the matmul
    for ix in left:
        if right.pop(ix, False):
            if ix in out:
                bat_inds.append(ix)
            else:
                con_inds.append(ix)
        elif ix in out:
            a_keep.append(ix)
    # now only indices unique to right remain
    for ix in right:
        if ix in out:
            b_keep.append(ix)

    if not con_inds:
        # contraction is pure multiplication, prepare inputs differently
        return _parse_eq_to_pure_multiplication(
            a_term, shape_a, b_term, shape_b, out
        )

    # only need the size one indices that appear in the output
    singletons = [ix for ix in out if ix in singletons]

    # take diagonal, remove any trivial axes and transpose left
    desired_a = "".join((*bat_inds, *a_keep, *con_inds))
    if a_term != desired_a:
        eq_a = f"{a_term}->{desired_a}"
    else:
        eq_a = None

    # take diagonal, remove any trivial axes and transpose right
    desired_b = "".join((*bat_inds, *con_inds, *b_keep))
    if b_term != desired_b:
        eq_b = f"{b_term}->{desired_b}"
    else:
        eq_b = None

    # then we want to reshape