backgrounds.shape assertion failure in `rasterize_to_pixels()` when `packed=True` (new to 1.5.3)

[eguendelman]

In gsplat 1.5.3, calling rasterization() with backgrounds set, we are getting an assertion failure in https://github.com/nerfstudio-project/gsplat/blob/main/gsplat/cuda/_wrapper.py#L598 which did not used to happen. And it only happens if packed=True (so setting packed=False is a workaround for now).

And we found that changing https://github.com/nerfstudio-project/gsplat/blob/main/gsplat/cuda/_wrapper.py#L582 to

image_dims = isect_offsets.shape[:-2]

fixes the issue, and also seems to better match the C++ code here https://github.com/nerfstudio-project/gsplat/blob/main/gsplat/cuda/csrc/Rasterization.cpp#L51

(There are a few other code locations, e.g. rasterize_to_pixels_2dgs(), that may also need to be fixed, though we do not use those code paths)

Here is a minimal repro for the issue:

import gsplat
import torch

# single scene with a single splat
means = torch.tensor([[0.0, 0.0, 0.0]]).cuda()
quats = torch.tensor([[1.0, 0.0, 0.0, 0.0]]).cuda()
scales = torch.tensor([[1.0, 1.0, 1.0]]).cuda()
opacities = torch.tensor([0.8]).cuda()
colors = torch.tensor([[0.8, 0.2, 0.2]]).cuda()

# single camera
viewmat = torch.eye(4).cuda()
viewmat[2, 3] = 4
width, height = 512, 512
K = torch.tensor([[width, 1.0, width/2], [1, height, height/2], [0, 0, 1]]).cuda()
background = torch.tensor([0.7, 0.9, 0.9]).cuda()

im, _, _ = gsplat.rasterization(means, quats, scales, opacities, colors,
                                viewmat.unsqueeze(0), K.unsqueeze(0), width, height,
                                backgrounds=background.unsqueeze(0), packed=True)

[zerolover]

see https://github.com/nerfstudio-project/gsplat/pull/742