Skip to content

Close incorrect quadratic-memory fix for u_sf_mRNA_factors - #445

Closed
TTTPOB wants to merge 2 commits into
BayraktarLab:masterfrom
TTTPOB:fix/deterministic-event-dim
Closed

Close incorrect quadratic-memory fix for u_sf_mRNA_factors#445
TTTPOB wants to merge 2 commits into
BayraktarLab:masterfrom
TTTPOB:fix/deterministic-event-dim

Conversation

@TTTPOB

@TTTPOB TTTPOB commented Jul 10, 2026

Copy link
Copy Markdown

Correction

This PR is being closed because its memory justification was incorrect. The (B, 1, B, F) Delta.v is a zero-stride expanded view that shares the original (B, F) mRNA storage; its logical numel() does not represent physical allocation.

pyro.deterministic() uses .mask(False), and Pyro's MaskedDistribution.log_prob() returns an expanded scalar zero without evaluating or materializing the underlying Delta value. The inspected TraceELBO path only consumes the resulting scalar log_prob_sum.

Moreover, event_dim=0 gives the Delta a (B, F) batch shape, which causes Delta.__init__ to allocate a real (B, F) log_density tensor. The proposed change therefore does not eliminate quadratic physical storage and can add a linear allocation.

The regression tests correctly identify a change in logical distribution shape, but they do not demonstrate reduced VRAM. No Cell2location code change is justified by the evidence presented here.

Detailed correction and measurements: #444

@TTTPOB TTTPOB changed the title Fix quadratic u_sf_mRNA_factors allocation Close incorrect quadratic-memory fix for u_sf_mRNA_factors Jul 10, 2026
@TTTPOB

TTTPOB commented Jul 10, 2026

Copy link
Copy Markdown
Author

Closing after correcting the allocation analysis in #444. The expanded Delta is a zero-stride view rather than quadratic physical storage, and event_dim=0 can add a linear log_density allocation.

@TTTPOB TTTPOB closed this Jul 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant