Preserve zero-length dimensions during broadcast shape inference

[Robertboy18]

Description

Hi! I was looking at einsum broadcasting semantics while formalizing some tensor shape behavior in Lean/TorchLean and noticed a discrepancy around zero-length dimensions.

find_output_shape currently uses the largest dimension for each output index. For a valid broadcast between dimensions 0 and 1, that infers 1, even though the broadcast result has dimension 0. In a multi-operand contraction this incorrect intermediate shape can select a contraction that raises a shape mismatch, while numpy.einsum evaluates the same expression successfully.

This change preserves the first non-singleton dimension, including 0, and adds parser-level and end-to-end regression coverage.

I also have the corresponding small TorchLean/Lean formalization of the 0/1 broadcast case and can attach it if that would be useful :)

Tests

• pytest -q opt_einsum/tests/test_parser.py opt_einsum/tests/test_edge_cases.py opt_einsum/tests/test_contract.py opt_einsum/tests/test_blas.py (7552 passed)
• uv run --extra test pytest -q (167 passed, 120 skipped)
• Ruff and mypy

Status

• Ready to go

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 89.88%. Comparing base (2ebc7b8) to head (670fd2b).

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.