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

Code before:
@functools.lru_cache(2**12)
def _parse_eq_to_batch_matmul(eq, shape_a, shape_b):
    """Cached parsing of a two term einsum equation into the necessary
    sequence of arguments for contracttion via batched matrix multiplication.
    The steps we need to specify are:

        1. Remove repeated and trivial indices from the left and right terms,
           and transpose them, done as a single einsum.
        2. Fuse the remaining indices so we have two 3D tensors.
        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)


Review comment (suggestion):
```suggestion
    left = {}
```

there's a ruff rule for this I believe, although it might not be enabled (yet)

Code after:
@functools.lru_cache(2**12)
def _parse_eq_to_batch_matmul(eq, shape_a, shape_b):
    """Cached parsing of a two term einsum equation into the necessary
    sequence of arguments for contracttion via batched matrix multiplication.
    The steps we need to specify are:

        1. Remove repeated and trivial indices from the left and right terms,
           and transpose them, done as a single einsum.
        2. Fuse the remaining indices so we have two 3D tensors.
        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)