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

Code before:
        3. Perform the batched matrix multiplication.
        4. Unfuse the output to get the desired final index order.

    """
    lhs, out = eq.split("->")
    a_term, b_term = lhs.split(",")

    if len(a_term) != len(shape_a):
        raise ValueError(f"Term '{a_term}' does not match shape {shape_a}.")
    if len(b_term) != len(shape_b):
        raise ValueError(f"Term '{b_term}' does not match shape {shape_b}.")

    sizes = {}
    singletons = set()

    # parse left term to unique indices with size > 1
    left = dict()
    for ix, d in zip(a_term, shape_a):
        if d == 1:
            # everything (including broadcasting) works nicely if simply ignore
            # such dimensions, but we do need to track if they appear in output
            # and thus should be reintroduced later
            singletons.add(ix)
            continue
        if sizes.setdefault(ix, d) != d:
            # set or check size
            raise ValueError(
                f"Index {ix} has mismatched sizes {sizes[ix]} and {d}."
            )
        left[ix] = True

    # parse right term to unique indices with size > 1
    right = dict()
    for ix, d in zip(b_term, shape_b):
        # broadcast indices (size 1 on one input and size != 1
        # on the other) should not be treated as singletons
        if d == 1:
            if ix not in left:
                singletons.add(ix)
            continue
        singletons.discard(ix)

        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

Review comment (suggestion):
```suggestion
            # set and check size
```

Code after:
        3. Perform the batched matrix multiplication.
        4. Unfuse the output to get the desired final index order.

    """
    lhs, out = eq.split("->")
    a_term, b_term = lhs.split(",")

    if len(a_term) != len(shape_a):
        raise ValueError(f"Term '{a_term}' does not match shape {shape_a}.")
    if len(b_term) != len(shape_b):
        raise ValueError(f"Term '{b_term}' does not match shape {shape_b}.")

    sizes = {}
    singletons = set()

    # parse left term to unique indices with size > 1
    left = {}
    for ix, d in zip(a_term, shape_a):
        if d == 1:
            # everything (including broadcasting) works nicely if simply ignore
            # such dimensions, but we do need to track if they appear in output
            # and thus should be reintroduced later
            singletons.add(ix)
            continue
        if sizes.setdefault(ix, d) != d:
            # set and check size
            raise ValueError(
                f"Index {ix} has mismatched sizes {sizes[ix]} and {d}."
            )
        left[ix] = True

    # parse right term to unique indices with size > 1
    right = {}
    for ix, d in zip(b_term, shape_b):
        # broadcast indices (size 1 on one input and size != 1
        # on the other) should not be treated as singletons
        if d == 1:
            if ix not in left:
                singletons.add(ix)
            continue
        singletons.discard(ix)

        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