An incorrect VRAM allocation report caused by codex hallucination

[TTTPOB]

Correction

Important

The original report incorrectly treated Delta.v.numel() * element_size() as physically allocated memory. The expanded (B, 1, B, F) shape is real, but it is a zero-stride view and does not allocate O(B^2 * F) storage. I am correcting and closing this report because the proposed event_dim=0 change is not a memory fix.

What remains true

With the current call inside obs_plate(dim=-2):

pyro.deterministic("u_sf_mRNA_factors", mRNA)

the deterministic site's unwrapped Delta.v has logical shape (B, 1, B, F), while the returned trace value remains (B, F).

What the original report got wrong

Delta.expand() uses Tensor.expand(). The resulting Delta.v shares the original (B, F) mRNA storage and has a zero stride on the added plate dimension. Therefore its logical numel() is not its allocated storage size.

For B=16 and F=4:

Tensor	Shape	Stride	Logical bytes	Storage bytes
Delta.v	(16, 1, 16, 4)	(0, 64, 4, 1)	4,096	256, shared with mRNA
deterministic log_prob	(16, 1)	(0, 0)	64	4

For B=8192 and F=9, the previously reported 2.25 GiB is only the logical size of the expanded view. Its physical backing storage is 8192 * 9 * 4 = 294,912 bytes (0.28125 MiB), shared with the original mRNA tensor.

Why ELBO evaluation does not materialize the view

pyro.deterministic() wraps the Delta with .mask(False). In Pyro 1.9.1, MaskedDistribution.log_prob() short-circuits this case:

if self._mask is False:
    shape = broadcast_shape(
        self.base_dist.batch_shape,
        value.shape[: value.dim() - self.event_dim],
    )
    return torch.zeros((), device=value.device).expand(shape)

The returned zeros are themselves an expanded scalar view. Trace.compute_log_prob() stores that view and its scalar sum; the inspected TraceELBO path does not call contiguous() or clone() on the expanded deterministic value.

Why event_dim=0 is not a memory optimization

Setting event_dim=0 changes the logical Delta.v shape to (B, F), but it does not reduce the shared mRNA backing storage. It also changes the Delta batch shape to (B, F), causing Delta.__init__ to allocate a real (B, F) log_density tensor. Under the old default, log_density is a scalar expanded view.

Measured at B=128, F=4:

Behavior	Delta.v storage	log_density storage
default event_dim=2	2,048 bytes, shared	4 bytes
proposed event_dim=0	2,048 bytes, shared	2,048 bytes

The proposed change therefore does not remove a quadratic physical allocation and can add one extra linear (B, F) allocation.

Conclusion

The trace contains an unusual expanded logical shape, but the evidence in this report does not demonstrate a quadratic CPU/GPU memory allocation in normal training or posterior collection. The memory-based justification for PR #445 is invalid, so I am closing both the PR and this issue. I apologize for the incorrect allocation claim.

Reproduction archive

The archive is retained for transparency. Its shape assertions remain valid, but fields calculated as numel() * element_size() describe logical tensor size, not allocated storage.

cell2location-deterministic-repro.zip

[TTTPOB]

Closing this report after correcting the distinction between logical tensor shape and physical storage. The attached archive remains available for transparency.

I was investigating some large vram occupation and codex leds me there. it looks tenable so i posted here. But after some more review I found it is not the cause, and codex is wrong here.

Sorry here, should have reviewed more carefully before posting it here.