⚡ Bolt: [Performance] Optimize einsum broadcasting overhead in MMLE expected counts

[seonghobae]

Optimization hypothesis

This Draft replaces two NumPy reference/fallback expected-count contractions in python/fast_mlsirm/estimators/marginal.py with context- and dimension-partitioned 2-D matrix products. The algebraic direction may avoid advanced-indexed 4-D contraction work and may use the configured BLAS implementation.

Current exact identity

• Protected base: main at 4d910ed650f384ff882c8b5fba6a8b08fd532236.
• Current head: 44f172c0c08795bc3670eecd32df6f7ba7c14326.
• Exact-head CI, Security Scan, and SAST passed, but no changed-path regression test, allocation evidence, reproducible benchmark manifest, or current-head formal review exists.
• CodeRabbit's automatic review was rate-limited; that is not passing review evidence.

Scientific and architecture boundary

• Production psychometric arithmetic remains Rust-first. This change is confined to the retained NumPy reference/fallback estimator and must preserve equation-level parity with the Rust-owned production path.
• Runtime improvement is shape-, missingness-, context-balance-, layout-, NumPy-, BLAS-, thread-, and hardware-dependent. The current 90x+ statement and 1.5 s -> 0.02 s journal claim are not accepted without a reproducible benchmark artifact and must not survive to merge as universal product evidence.
• Passing the pre-existing suite is insufficient because this head changes numerical contraction ownership without adding a focused regression.

Dependency order

PR #564 is the earlier active branch that also changes python/fast_mlsirm/estimators/marginal.py. Do not stack, rebase, or dispatch a competing source writer while #564 has a live writer lease. After #564 reaches protected main, refresh this branch from the exact protected tip and reevaluate whether the optimization still applies before retaining any implementation.

Required test-first completion

Before Ready status, establish focused RED -> GREEN evidence that independently reconstructs the previous einsum equations and compares them with the proposed contractions for:

• single and multiple contexts, including an unused context;
• one and multiple factor dimensions, including a dimension with no assigned item;
• complete and partially missing responses;
• balanced and highly unbalanced context membership;
• ordinary and non-contiguous input layouts where the public contract permits them;
• expected person log-likelihood tensors, nbar, rbar, and mbar at tight declared tolerances;
• input-order and context-label permutations; and
• bounded peak-memory/allocation behavior at realistic large shapes.

Use an environment-recorded benchmark with hardware, OS, Python, NumPy, BLAS/thread configuration, dtype/layout, warm-up, repetitions, elapsed distribution, and peak memory. Report observations rather than a universal speedup. Remove unrelated whole-file formatting churn, replace the temporary .jules/bolt.md assertion with authoritative changelog/APA 7 doctoring when the contract is proven, and rerun exact-head Python/Rust/PyO3/package/GPU-no-skip/fuzz/security/SAST gates.

Keep Draft until the predecessor is integrated, the refreshed exact head has the bounded tests and evidence above, no valid review finding remains, and repository merge policy is satisfied.

---

Originally created by Jules task 4768818120839347618; subsequent governance state is maintained by the repository writer loop.

[google-labs-jules]

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.

---

For security, I will only act on instructions from the user who triggered this task.

[coderabbitai]

Warning

Review limit reached

@seonghobae, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 56 minutes

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: a28fc8b2-1758-4586-9742-ff471e769546

📥 Commits

Reviewing files that changed from the base of the PR and between 4d910ed and 44f172c.

📒 Files selected for processing (2)

• .jules/bolt.md
• python/fast_mlsirm/estimators/marginal.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[seonghobae]

Closing without merge rather than retaining a knowingly non-integrable speculative branch behind #564. Exact head 44f172c0c08795bc3670eecd32df6f7ba7c14326 changes the same active python/fast_mlsirm/estimators/marginal.py boundary as #564, adds 248/82 lines dominated by whole-file formatter churn, commits a non-authoritative .jules/bolt.md entry with an incorrect 2024 date, and embeds an unsupported 90x+/1.5s -> 0.02s claim. It has no focused equation-parity, missingness, layout, permutation, allocation, or reproducible benchmark evidence. The branch therefore carries unique but unaccepted experimental work, not integrated product value. After #564 reaches protected main, reevaluate the optimization from that exact tip and reintroduce only a minimal test-first diff if profiling still proves it material. No checks, reviews, or benchmark claims transfer.

[google-labs-jules]

Closing without merge rather than retaining a knowingly non-integrable speculative branch behind #564. Exact head 44f172c0c08795bc3670eecd32df6f7ba7c14326 changes the same active python/fast_mlsirm/estimators/marginal.py boundary as #564, adds 248/82 lines dominated by whole-file formatter churn, commits a non-authoritative .jules/bolt.md entry with an incorrect 2024 date, and embeds an unsupported 90x+/1.5s -> 0.02s claim. It has no focused equation-parity, missingness, layout, permutation, allocation, or reproducible benchmark evidence. The branch therefore carries unique but unaccepted experimental work, not integrated product value. After #564 reaches protected main, reevaluate the optimization from that exact tip and reintroduce only a minimal test-first diff if profiling still proves it material. No checks, reviews, or benchmark claims transfer.

Understood. Acknowledging that this work is now obsolete and stopping work on this task.