Skip to content

Count broadcast batches in optimized einsum handlers - #162

Open
AryanRG wants to merge 1 commit into
facebookresearch:mainfrom
AryanRG:fix/einsum-broadcast-flops
Open

AryanRG wants to merge 1 commit into
facebookresearch:mainfrom
AryanRG:fix/einsum-broadcast-flops

Conversation

@AryanRG

@AryanRG AryanRG commented Oct 28, 2026

Copy link
Copy Markdown

The two optimized einsum handlers use the left input's batch size, which undercounts FLOPs when that dimension broadcasts. For example, a batch of four is reported as 30 FLOPs instead of 120.

This uses the traced output batch size and adds broadcast and empty-batch regressions. The Python suite passes with CUDA tests skipped. Existing repository-wide lint failures reproduce on the base branch; the changed code passes its checks.

Prepared with AI assistance.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant