Improve docs: The get_mlp_output method expects enc_feats with shape [B, R, S, F] or [B, H, W, S, F] (as documented in lines 599-604), but get_gaussian_affinity_output passes enc_feats with shape [N, F] after the rgb_to_sh transformation. This shape mismatch will likely cause runtime errors or incorrect behavior. Consider reshaping enc_feats appropriately or verifying that get_mlp_output can handle this alternate shape.

[swahtz]

The get_mlp_output method expects enc_feats with shape [B, R, S, F] or [B, H, W, S, F] (as documented in lines 599-604), but get_gaussian_affinity_output passes enc_feats with shape [N, F] after the rgb_to_sh transformation. This shape mismatch will likely cause runtime errors or incorrect behavior. Consider reshaping enc_feats appropriately or verifying that get_mlp_output can handle this alternate shape.

            # Convert per-gaussian RGB features to spherical harmonics coefficients
            enc_feats = rgb_to_sh(enc_feats)
        else:
            with nvtx.range("update_gs_features"):
                # Update sh0 with current gs_features and reuse the model
                self._gs_model_for_render.sh0 = self.gs_features
                gs3d_enc_feats = self._gs_model_for_render
            # Use the current Gaussian features directly
            enc_feats = self.gs_features

        epsilon = 1e-6
        enc_feats = enc_feats / (torch.linalg.norm(enc_feats, dim=-1, keepdim=True) + epsilon)

        # Reshape enc_feats to match get_mlp_output expected shape [B, R, S, F]
        enc_feats_in = enc_feats.unsqueeze(0).unsqueeze(2)  # [1, N, 1, F]

        # Apply MLP
        gfeats = self.get_mlp_output(enc_feats_in, scale)
        # Remove artificial batch and sample dimensions to return per-gaussian features
        gfeats = gfeats.squeeze(0).squeeze(1)

Originally posted by @Copilot in #49 (comment)