"""Cross-view evidence construction and bounded residual completion in a typed 3D state.""" from __future__ import annotations import math from collections.abc import Sequence import torch import torch.nn.functional as functional from torch import Tensor, nn from mcss.geometry import make_voxel_centers, project_world from mcss.model.encoder2d import PyramidImageEncoder from mcss.model.state_encoder import Residual3DBlock from mcss.types import ( Cameras, DualEvidence, SceneState, StateAppearance, StateEvidence, TransportEvidence, ) class EvidenceResidualStateEncoder(nn.Module): """Build an inspectable evidence state and add only a bounded, gated completion residual.""" def __init__( self, *, voxel_resolution: int | Sequence[int], image_feature_dim: int, state_feature_dim: int, refinement_blocks: int, evidence_temperature: float, observed_residual_floor: float, completion_residual_scale: float, appearance_resolution_scale: int = 1, ) -> None: super().__init__() if isinstance(voxel_resolution, int): resolution = (voxel_resolution,) * 3 else: resolution = tuple(voxel_resolution) if len(resolution) != 3 or min(resolution) < 4: raise ValueError("voxel_resolution must contain D, H, W values of at least four") if image_feature_dim < 4 or state_feature_dim < 4: raise ValueError("feature dimensions must be at least four") if refinement_blocks < 0: raise ValueError("refinement_blocks cannot be negative") if evidence_temperature <= 0: raise ValueError("evidence_temperature must be positive") if not 0.0 <= observed_residual_floor <= 1.0: raise ValueError("observed_residual_floor must be within [0, 1]") if completion_residual_scale <= 0: raise ValueError("completion_residual_scale must be positive") if appearance_resolution_scale not in {1, 2}: raise ValueError("appearance_resolution_scale must be one or two") self.voxel_resolution = resolution self.evidence_temperature = float(evidence_temperature) self.observed_residual_floor = float(observed_residual_floor) self.completion_residual_scale = float(completion_residual_scale) self.appearance_resolution_scale = appearance_resolution_scale self.image_encoder = PyramidImageEncoder(image_feature_dim) observation_channels = image_feature_dim + 3 + 3 + 1 fused_channels = observation_channels * 2 + 3 hidden_channels = max(32, state_feature_dim * 2) self.input_projection = nn.Sequential( nn.Conv3d(fused_channels, hidden_channels, 3, padding=1, bias=False), nn.GroupNorm(_group_count(hidden_channels), hidden_channels), nn.SiLU(inplace=True), ) self.refinement = nn.Sequential( *(Residual3DBlock(hidden_channels) for _ in range(refinement_blocks)) ) self.density_residual_head = nn.Conv3d(hidden_channels, 1, 1) self.color_residual_head = nn.Conv3d(hidden_channels, 3, 1) self.variance_head = nn.Conv3d(hidden_channels, 1, 1) self.feature_head = nn.Conv3d(hidden_channels, state_feature_dim, 1) nn.init.zeros_(self.density_residual_head.weight) nn.init.zeros_(self.density_residual_head.bias) nn.init.zeros_(self.color_residual_head.weight) nn.init.zeros_(self.color_residual_head.bias) nn.init.constant_(self.variance_head.bias, -3.0) self.appearance_input_projection: nn.Module | None = None self.appearance_refinement: nn.Module | None = None self.appearance_residual_head: nn.Conv3d | None = None if appearance_resolution_scale == 2: appearance_hidden_channels = max(16, state_feature_dim) self.appearance_input_projection = nn.Sequential( nn.Conv3d( hidden_channels + 5, appearance_hidden_channels, 3, padding=1, bias=False, ), nn.GroupNorm( _group_count(appearance_hidden_channels), appearance_hidden_channels ), nn.SiLU(inplace=True), ) self.appearance_refinement = nn.Sequential( *( Residual3DBlock(appearance_hidden_channels) for _ in range(max(1, refinement_blocks // 2)) ) ) self.appearance_residual_head = nn.Conv3d(appearance_hidden_channels, 3, 1) nn.init.zeros_(self.appearance_residual_head.weight) nn.init.zeros_(self.appearance_residual_head.bias) def forward(self, context_rgb: Tensor, cameras: Cameras, bounds: Tensor) -> SceneState: if context_rgb.ndim != 5 or context_rgb.shape[2] != 3: raise ValueError("context_rgb must have shape [B, V, 3, H, W]") batch_size, view_count, _, image_height, image_width = context_rgb.shape if cameras.leading_shape != (batch_size, view_count): raise ValueError("context cameras must align with context_rgb") if cameras.image_size != (image_height, image_width): raise ValueError("context camera image size must match context_rgb") if bounds.shape != (batch_size, 2, 3): raise ValueError("bounds must have shape [B, 2, 3]") encoded = self.image_encoder( context_rgb.reshape(batch_size * view_count, 3, image_height, image_width) ) _, channels, feature_height, feature_width = encoded.shape encoded = encoded.reshape(batch_size, view_count, channels, feature_height, feature_width) points_per_view, normalized_grid, camera_depth, valid = _project_voxel_grid( bounds, cameras, self.voxel_resolution, image_height, image_width, ) sampled_features = _sample_images(encoded, normalized_grid) sampled_rgb = _sample_images(context_rgb, normalized_grid) evidence_features = _fixed_rgb_pyramid(context_rgb, normalized_grid) confidence, provenance = _cross_view_evidence( evidence_features, sampled_rgb, valid, temperature=self.evidence_temperature, ) unknown_probability = 1.0 - confidence camera_centers = cameras.c2w[..., :3, 3].unsqueeze(-2) view_direction = _normalize(points_per_view - camera_centers) scene_diagonal = torch.linalg.vector_norm(bounds[:, 1] - bounds[:, 0], dim=-1) normalized_depth = camera_depth / scene_diagonal[:, None, None].clamp_min(1e-6) observations = torch.cat( (sampled_features, sampled_rgb, view_direction, normalized_depth.unsqueeze(-1)), dim=-1 ) mean = (observations * provenance).sum(dim=1) variance = ((observations - mean[:, None]).square() * provenance).sum(dim=1) coverage = valid.to(observations.dtype).mean(dim=1, keepdim=False).unsqueeze(-1) fused = torch.cat((mean, variance, coverage, confidence, unknown_probability), dim=-1) fused_grid = _points_to_volume(fused, self.voxel_resolution) hidden = self.refinement(self.input_projection(fused_grid)) density_residual = ( torch.tanh(self.density_residual_head(hidden)) * self.completion_residual_scale ) color_logit_residual = ( torch.tanh(self.color_residual_head(hidden)) * self.completion_residual_scale ) confidence_grid = _points_to_volume(confidence, self.voxel_resolution) unknown_grid = 1.0 - confidence_grid provenance_grid = provenance.squeeze(-1).reshape( batch_size, view_count, *self.voxel_resolution ) base_color_points = _base_color_from_provenance(sampled_rgb, provenance) base_color = _points_to_volume(base_color_points, self.voxel_resolution).clamp(0.0, 1.0) base_density_logits = 4.0 * confidence_grid - 2.0 completion_gate = self.observed_residual_floor + ( 1.0 - self.observed_residual_floor ) * unknown_grid density_logits = base_density_logits + completion_gate * density_residual base_color_logits = torch.logit(base_color.clamp(1e-4, 1.0 - 1e-4)) color = torch.sigmoid(base_color_logits + completion_gate * color_logit_residual) log_variance = (self.variance_head(hidden) + 2.0 * unknown_grid).clamp(-8.0, 4.0) features = self.feature_head(hidden) evidence = StateEvidence( confidence=confidence_grid, unknown_probability=unknown_grid, completion_gate=completion_gate, provenance=provenance_grid, base_density_logits=base_density_logits, base_color=base_color, density_residual=density_residual, color_logit_residual=color_logit_residual, ) appearance = None if self.appearance_resolution_scale == 2: appearance = self._build_highres_appearance( context_rgb, cameras, bounds, hidden, ) return SceneState( density_logits, color, log_variance, bounds, features, evidence, appearance, ) def _build_highres_appearance( self, context_rgb: Tensor, cameras: Cameras, bounds: Tensor, hidden: Tensor, ) -> StateAppearance: if ( self.appearance_input_projection is None or self.appearance_refinement is None or self.appearance_residual_head is None ): raise RuntimeError("high-resolution appearance modules are not initialized") image_height, image_width = context_rgb.shape[-2:] resolution = tuple( axis * self.appearance_resolution_scale for axis in self.voxel_resolution ) _, normalized_grid, _, valid = _project_voxel_grid( bounds, cameras, resolution, image_height, image_width, ) sampled_rgb = _sample_images(context_rgb, normalized_grid) fixed_features = _fixed_rgb_pyramid(context_rgb, normalized_grid) confidence, provenance = _cross_view_evidence( fixed_features, sampled_rgb, valid, temperature=self.evidence_temperature, ) confidence_grid = _points_to_volume(confidence, resolution) unknown_grid = 1.0 - confidence_grid completion_gate = self.observed_residual_floor + ( 1.0 - self.observed_residual_floor ) * unknown_grid provenance_grid = provenance.squeeze(-1).reshape( context_rgb.shape[0], context_rgb.shape[1], *resolution ) base_color_points = _base_color_from_provenance(sampled_rgb, provenance) base_color = _points_to_volume(base_color_points, resolution).clamp(0.0, 1.0) lifted_hidden = functional.interpolate( hidden, size=resolution, mode="trilinear", align_corners=False, ) appearance_inputs = torch.cat( (lifted_hidden, base_color, confidence_grid, unknown_grid), dim=1 ) appearance_hidden = self.appearance_refinement( self.appearance_input_projection(appearance_inputs) ) residual = ( torch.tanh(self.appearance_residual_head(appearance_hidden)) * self.completion_residual_scale ) return StateAppearance( confidence=confidence_grid, unknown_probability=unknown_grid, completion_gate=completion_gate, provenance=provenance_grid, base_color=base_color, color_logit_residual=residual, ) class DualEvidenceStateEncoder(EvidenceResidualStateEncoder): """Add a detached local ray-competition gate without changing V5 trainable parameters.""" def __init__( self, *, voxel_resolution: int | Sequence[int], image_feature_dim: int, state_feature_dim: int, refinement_blocks: int, evidence_temperature: float, surface_peak_temperature: float, observed_residual_floor: float, completion_residual_scale: float, appearance_resolution_scale: int = 1, surface_peak_chunk_size: int = 4096, ) -> None: super().__init__( voxel_resolution=voxel_resolution, image_feature_dim=image_feature_dim, state_feature_dim=state_feature_dim, refinement_blocks=refinement_blocks, evidence_temperature=evidence_temperature, observed_residual_floor=observed_residual_floor, completion_residual_scale=completion_residual_scale, appearance_resolution_scale=appearance_resolution_scale, ) if surface_peak_temperature <= 0: raise ValueError("surface_peak_temperature must be positive") if surface_peak_chunk_size < 1: raise ValueError("surface_peak_chunk_size must be positive") self.surface_peak_temperature = float(surface_peak_temperature) self.surface_peak_chunk_size = int(surface_peak_chunk_size) def forward(self, context_rgb: Tensor, cameras: Cameras, bounds: Tensor) -> SceneState: v5_state = super().forward(context_rgb, cameras, bounds) if v5_state.evidence is None: raise RuntimeError("dual evidence requires the V5 evidence construction") appearance_evidence = v5_state.evidence peakness = _surface_peakness( context_rgb, cameras, bounds, self.voxel_resolution, evidence_temperature=self.evidence_temperature, surface_peak_temperature=self.surface_peak_temperature, chunk_size=self.surface_peak_chunk_size, ).to(dtype=appearance_evidence.confidence.dtype) appearance_confidence = appearance_evidence.confidence localization_support = (appearance_confidence * peakness).clamp(0.0, 1.0) localization_unknown = 1.0 - localization_support appearance_unknown = 1.0 - appearance_confidence localization_gate = self.observed_residual_floor + ( 1.0 - self.observed_residual_floor ) * localization_unknown appearance_gate = appearance_evidence.completion_gate base_density_logits = 4.0 * localization_support - 2.0 density_logits = ( base_density_logits + localization_gate * appearance_evidence.density_residual ) dual_evidence = DualEvidence( surface_localization_support=localization_support, appearance_confidence=appearance_confidence, surface_peakness=peakness, localization_unknown_probability=localization_unknown, appearance_unknown_probability=appearance_unknown, localization_completion_gate=localization_gate, appearance_completion_gate=appearance_gate, provenance=appearance_evidence.provenance, base_density_logits=base_density_logits, base_color=appearance_evidence.base_color, density_residual=appearance_evidence.density_residual, color_logit_residual=appearance_evidence.color_logit_residual, ) return SceneState( density_logits=density_logits, color=v5_state.color, log_variance=v5_state.log_variance, bounds=v5_state.bounds, features=v5_state.features, evidence=None, appearance=v5_state.appearance, dual_evidence=dual_evidence, ) class TransportEvidenceStateEncoder(EvidenceResidualStateEncoder): """Transport fixed local ray evidence to its metric 3D candidate positions.""" def __init__( self, *, voxel_resolution: int | Sequence[int], image_feature_dim: int, state_feature_dim: int, refinement_blocks: int, evidence_temperature: float, surface_peak_temperature: float, observed_residual_floor: float, completion_residual_scale: float, appearance_resolution_scale: int = 1, surface_peak_chunk_size: int = 4096, ) -> None: super().__init__( voxel_resolution=voxel_resolution, image_feature_dim=image_feature_dim, state_feature_dim=state_feature_dim, refinement_blocks=refinement_blocks, evidence_temperature=evidence_temperature, observed_residual_floor=observed_residual_floor, completion_residual_scale=completion_residual_scale, appearance_resolution_scale=appearance_resolution_scale, ) if surface_peak_temperature <= 0: raise ValueError("surface_peak_temperature must be positive") if surface_peak_chunk_size < 1: raise ValueError("surface_peak_chunk_size must be positive") self.surface_peak_temperature = float(surface_peak_temperature) self.surface_peak_chunk_size = int(surface_peak_chunk_size) def forward(self, context_rgb: Tensor, cameras: Cameras, bounds: Tensor) -> SceneState: v5_state = super().forward(context_rgb, cameras, bounds) if v5_state.evidence is None: raise RuntimeError("transport evidence requires the V5 evidence construction") appearance_evidence = v5_state.evidence coverage, transported_appearance, peakness, localization_support = ( _transport_surface_evidence( context_rgb, cameras, bounds, self.voxel_resolution, appearance_evidence.confidence, evidence_temperature=self.evidence_temperature, surface_peak_temperature=self.surface_peak_temperature, chunk_size=self.surface_peak_chunk_size, ) ) evidence_dtype = appearance_evidence.confidence.dtype coverage = coverage.to(dtype=evidence_dtype) transported_appearance = transported_appearance.to(dtype=evidence_dtype) peakness = peakness.to(dtype=evidence_dtype) localization_support = localization_support.to(dtype=evidence_dtype) appearance_confidence = appearance_evidence.confidence localization_unknown = 1.0 - localization_support appearance_unknown = 1.0 - appearance_confidence localization_gate = self.observed_residual_floor + ( 1.0 - self.observed_residual_floor ) * localization_unknown appearance_gate = appearance_evidence.completion_gate base_density_logits = 4.0 * localization_support - 2.0 density_logits = ( base_density_logits + localization_gate * appearance_evidence.density_residual ) transport_evidence = TransportEvidence( surface_localization_support=localization_support, transported_appearance_confidence=transported_appearance, surface_peakness=peakness, transport_coverage=coverage, appearance_confidence=appearance_confidence, localization_unknown_probability=localization_unknown, appearance_unknown_probability=appearance_unknown, localization_completion_gate=localization_gate, appearance_completion_gate=appearance_gate, provenance=appearance_evidence.provenance, base_density_logits=base_density_logits, base_color=appearance_evidence.base_color, density_residual=appearance_evidence.density_residual, color_logit_residual=appearance_evidence.color_logit_residual, ) return SceneState( density_logits=density_logits, color=v5_state.color, log_variance=v5_state.log_variance, bounds=v5_state.bounds, features=v5_state.features, evidence=None, appearance=v5_state.appearance, dual_evidence=None, transport_evidence=transport_evidence, ) @torch.no_grad() def _trilinear_splat_candidates( points: Tensor, values: Tensor, valid: Tensor, bounds: Tensor, resolution: tuple[int, int, int], ) -> Tensor: """Splat candidate values onto native voxel centers with boundary mass conservation.""" if points.ndim != 4 or points.shape[-1] != 3: raise ValueError("points must have shape [B, C, K, 3]") candidate_shape = points.shape[:-1] scalar_values = values.ndim == 3 and values.shape == candidate_shape channel_values = ( values.ndim == 4 and values.shape[:-1] == candidate_shape and values.shape[-1] >= 1 ) if not scalar_values and not channel_values: raise ValueError( "values must have shape [B, C, K] or [B, C, K, M] aligned with points" ) if valid.shape != candidate_shape: raise ValueError("valid must align with candidate points") if valid.dtype != torch.bool: raise ValueError("valid must be boolean") if bounds.shape != (points.shape[0], 2, 3): raise ValueError("bounds must have shape [B, 2, 3]") if len(resolution) != 3 or min(resolution) < 1: raise ValueError("resolution must contain positive D, H, W values") if not torch.isfinite(points).all() or not torch.isfinite(values).all(): raise ValueError("splat inputs must be finite") points = points.detach().float() values = values.detach().float() bounds = bounds.detach().float() depth, height, width = resolution axis_sizes = torch.tensor( [width, height, depth], device=points.device, dtype=points.dtype, ) cell_size = (bounds[:, 1] - bounds[:, 0]) / axis_sizes coordinates = ( (points - bounds[:, None, None, 0]) / cell_size[:, None, None] - 0.5 ) base = torch.floor(coordinates).to(torch.int64) fraction = coordinates - base.to(coordinates.dtype) neighbor_offsets = torch.tensor( [ [0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1], ], device=points.device, dtype=torch.int64, ) neighbor = base[..., None, :] + neighbor_offsets within_grid = ( (neighbor[..., 0] >= 0) & (neighbor[..., 0] < width) & (neighbor[..., 1] >= 0) & (neighbor[..., 1] < height) & (neighbor[..., 2] >= 0) & (neighbor[..., 2] < depth) ) axis_weights = torch.where( neighbor_offsets.to(torch.bool)[None, None, None], fraction[..., None, :], 1.0 - fraction[..., None, :], ) weights = axis_weights.prod(dim=-1) weights *= (within_grid & valid[..., None]).to(weights.dtype) weight_sum = weights.sum(dim=-1, keepdim=True) weights = torch.where(weight_sum > 0, weights / weight_sum.clamp_min(1e-12), weights) x = neighbor[..., 0].clamp(0, width - 1) y = neighbor[..., 1].clamp(0, height - 1) z = neighbor[..., 2].clamp(0, depth - 1) linear_index = z * (height * width) + y * width + x values_by_channel = values.unsqueeze(-1) if scalar_values else values channel_count = values_by_channel.shape[-1] output = torch.zeros( points.shape[0], channel_count, depth * height * width, device=points.device, dtype=torch.float32, ) indices = linear_index.reshape(points.shape[0], 1, -1).expand( -1, channel_count, -1 ) contributions = (weights[..., None] * values_by_channel[..., None, :]).permute( 0, 4, 1, 2, 3 ) output.scatter_add_( 2, indices, contributions.reshape(points.shape[0], channel_count, -1), ) return output[:, 0] if scalar_values else output @torch.no_grad() def _transport_surface_evidence( context_rgb: Tensor, cameras: Cameras, bounds: Tensor, resolution: tuple[int, int, int], appearance_confidence: Tensor, *, evidence_temperature: float, surface_peak_temperature: float, chunk_size: int, ) -> tuple[Tensor, Tensor, Tensor, Tensor]: """Return coverage, transported appearance, peakness, and localization support.""" if context_rgb.ndim != 5 or context_rgb.shape[2] != 3: raise ValueError("context_rgb must have shape [B, V, 3, H, W]") batch_size, view_count, _, image_height, image_width = context_rgb.shape if cameras.leading_shape != (batch_size, view_count): raise ValueError("context cameras must align with context_rgb") if cameras.image_size != (image_height, image_width): raise ValueError("context camera image size must match context_rgb") if bounds.shape != (batch_size, 2, 3): raise ValueError("bounds must have shape [B, 2, 3]") if appearance_confidence.shape != (batch_size, 1, *resolution): raise ValueError("appearance_confidence must align with the native grid") if evidence_temperature <= 0 or surface_peak_temperature <= 0: raise ValueError("evidence temperatures must be positive") if chunk_size < 1: raise ValueError("chunk_size must be positive") zero = torch.zeros( batch_size, 1, *resolution, device=context_rgb.device, dtype=torch.float32, ) if view_count < 3: return zero, zero.clone(), zero.clone(), zero.clone() with torch.autocast(device_type=context_rgb.device.type, enabled=False): rgb = context_rgb.detach().float() bounds_fp32 = bounds.detach().float() cameras_fp32 = cameras.to(dtype=torch.float32) appearance_points = appearance_confidence.detach().float().reshape(batch_size, -1) centers = make_voxel_centers(bounds_fp32, resolution).reshape(batch_size, -1, 3) point_count = centers.shape[1] center_points_per_view = centers[:, None].expand( batch_size, view_count, point_count, 3 ) center_pixels, _, _ = project_world(center_points_per_view, cameras_fp32) center_grid = _pixels_to_grid(center_pixels, image_height, image_width) reference_features = _fixed_rgb_pyramid(rgb, center_grid) reference_rgb = _sample_images(rgb, center_grid) depth, height, width = resolution axis_sizes = torch.tensor( [width, height, depth], device=bounds_fp32.device, dtype=bounds_fp32.dtype, ) delta = ((bounds_fp32[:, 1] - bounds_fp32[:, 0]) / axis_sizes).amin(dim=-1) offsets = torch.tensor( [-2.0, -1.0, 0.0, 1.0, 2.0], device=bounds_fp32.device, dtype=bounds_fp32.dtype, ) camera_centers = cameras_fp32.c2w[..., :3, 3] accumulators = torch.zeros( batch_size, 3, point_count, device=bounds_fp32.device, dtype=torch.float32, ) other_view_mask = ~torch.eye( view_count, device=bounds_fp32.device, dtype=torch.bool ) for reference_index in range(view_count): for start in range(0, point_count, chunk_size): stop = min(start + chunk_size, point_count) center_chunk = centers[:, start:stop] direction = _normalize( center_chunk - camera_centers[:, reference_index, None] ) candidates = center_chunk[:, :, None] + ( direction[:, :, None] * delta[:, None, None, None] * offsets[None, None, :, None] ) inside = ( (candidates >= bounds_fp32[:, None, None, 0]) & (candidates <= bounds_fp32[:, None, None, 1]) ).all(dim=-1) chunk_points = candidates.reshape(batch_size, -1, 3) points_per_view = chunk_points[:, None].expand( batch_size, view_count, chunk_points.shape[1], 3 ) pixels, _, projected_valid = project_world(points_per_view, cameras_fp32) normalized_grid = _pixels_to_grid(pixels, image_height, image_width) sampled_features = _fixed_rgb_pyramid(rgb, normalized_grid).reshape( batch_size, view_count, stop - start, 5, -1 ) sampled_rgb = _sample_images(rgb, normalized_grid).reshape( batch_size, view_count, stop - start, 5, 3 ) projected_valid = projected_valid.reshape( batch_size, view_count, stop - start, 5 ) observer_valid = projected_valid & inside[:, None] observer_valid &= other_view_mask[reference_index][None, :, None, None] observer_count = observer_valid.sum(dim=1) reference_valid = projected_valid[:, reference_index] & inside candidate_valid = reference_valid & (observer_count >= 2) reference_feature = reference_features[ :, reference_index, start:stop ][:, None, :, None] reference_color = reference_rgb[:, reference_index, start:stop][ :, None, :, None ] feature_agreement = ( functional.cosine_similarity( reference_feature, sampled_features, dim=-1, eps=1e-6, ) + 1.0 ) * 0.5 rgb_error = (sampled_rgb - reference_color).square().mean(dim=-1) rgb_agreement = torch.exp(-rgb_error / evidence_temperature) pair_score = feature_agreement * rgb_agreement score = ( pair_score * observer_valid.to(pair_score.dtype) ).sum(dim=1) / observer_count.to(pair_score.dtype).clamp_min(1.0) excess, contributes = _positive_candidate_excess( score, candidate_valid, temperature=surface_peak_temperature, ) transport_valid = candidate_valid & contributes[..., None] source_appearance = appearance_points[:, start:stop, None].expand_as(excess) values = torch.stack( ( torch.ones_like(excess), source_appearance, source_appearance * excess, ), dim=-1, ) accumulators += _trilinear_splat_candidates( candidates, values, transport_valid, bounds_fp32, resolution, ) coverage = accumulators[:, 0] appearance_mass = accumulators[:, 1] localized_mass = accumulators[:, 2] transported_appearance = torch.where( coverage > 0, appearance_mass / coverage.clamp_min(1e-12), torch.zeros_like(coverage), ).clamp(0.0, 1.0) peakness = torch.where( appearance_mass > 0, localized_mass / appearance_mass.clamp_min(1e-12), torch.zeros_like(appearance_mass), ).clamp(0.0, 1.0) localization_support = torch.where( coverage > 0, localized_mass / coverage.clamp_min(1e-12), torch.zeros_like(coverage), ).clamp(0.0, 1.0) def volume(values: Tensor) -> Tensor: return values.reshape(batch_size, 1, depth, height, width).detach() return ( volume(coverage), volume(transported_appearance), volume(peakness), volume(localization_support), ) @torch.no_grad() def _surface_peakness( context_rgb: Tensor, cameras: Cameras, bounds: Tensor, resolution: tuple[int, int, int], *, evidence_temperature: float, surface_peak_temperature: float, chunk_size: int, _telemetry_records: list[tuple[Tensor, Tensor, Tensor, Tensor | None]] | None = None, _telemetry_appearance_confidence: Tensor | None = None, ) -> Tensor: """Return native-grid local ray-competition support in detached FP32.""" if context_rgb.ndim != 5 or context_rgb.shape[2] != 3: raise ValueError("context_rgb must have shape [B, V, 3, H, W]") batch_size, view_count, _, image_height, image_width = context_rgb.shape if cameras.leading_shape != (batch_size, view_count): raise ValueError("context cameras must align with context_rgb") if bounds.shape != (batch_size, 2, 3): raise ValueError("bounds must have shape [B, 2, 3]") if _telemetry_appearance_confidence is not None and ( _telemetry_appearance_confidence.shape != (batch_size, 1, *resolution) ): raise ValueError("telemetry appearance confidence must align with the native grid") if evidence_temperature <= 0 or surface_peak_temperature <= 0: raise ValueError("evidence temperatures must be positive") if chunk_size < 1: raise ValueError("chunk_size must be positive") if view_count < 2: return torch.zeros( batch_size, 1, *resolution, device=context_rgb.device, dtype=torch.float32, ) with torch.autocast(device_type=context_rgb.device.type, enabled=False): rgb = context_rgb.detach().float() bounds_fp32 = bounds.detach().float() cameras_fp32 = cameras.to(dtype=torch.float32) appearance_points = ( None if _telemetry_appearance_confidence is None else _telemetry_appearance_confidence.detach().float().reshape(batch_size, -1) ) centers = make_voxel_centers(bounds_fp32, resolution).reshape(batch_size, -1, 3) point_count = centers.shape[1] center_points_per_view = centers[:, None].expand( batch_size, view_count, point_count, 3 ) center_pixels, _, center_valid = project_world( center_points_per_view, cameras_fp32 ) center_grid = _pixels_to_grid(center_pixels, image_height, image_width) reference_features = _fixed_rgb_pyramid(rgb, center_grid) reference_rgb = _sample_images(rgb, center_grid) depth, height, width = resolution axis_sizes = torch.tensor( [width, height, depth], device=bounds_fp32.device, dtype=bounds_fp32.dtype, ) delta = ((bounds_fp32[:, 1] - bounds_fp32[:, 0]) / axis_sizes).amin(dim=-1) offsets = torch.tensor( [-2.0, -1.0, 0.0, 1.0, 2.0], device=bounds_fp32.device, dtype=bounds_fp32.dtype, ) camera_centers = cameras_fp32.c2w[..., :3, 3] peakness_sum = torch.zeros( batch_size, point_count, device=bounds_fp32.device, dtype=torch.float32 ) reference_count = torch.zeros_like(peakness_sum) other_view_mask = ~torch.eye( view_count, device=bounds_fp32.device, dtype=torch.bool ) for reference_index in range(view_count): for start in range(0, point_count, chunk_size): stop = min(start + chunk_size, point_count) center_chunk = centers[:, start:stop] direction = _normalize( center_chunk - camera_centers[:, reference_index, None] ) candidates = center_chunk[:, :, None] + ( direction[:, :, None] * delta[:, None, None, None] * offsets[None, None, :, None] ) inside = ( (candidates >= bounds_fp32[:, None, None, 0]) & (candidates <= bounds_fp32[:, None, None, 1]) ).all(dim=-1) chunk_points = candidates.reshape(batch_size, -1, 3) points_per_view = chunk_points[:, None].expand( batch_size, view_count, chunk_points.shape[1], 3 ) pixels, _, projected_valid = project_world( points_per_view, cameras_fp32 ) normalized_grid = _pixels_to_grid(pixels, image_height, image_width) sampled_features = _fixed_rgb_pyramid(rgb, normalized_grid).reshape( batch_size, view_count, stop - start, 5, -1 ) sampled_rgb = _sample_images(rgb, normalized_grid).reshape( batch_size, view_count, stop - start, 5, 3 ) projected_valid = projected_valid.reshape( batch_size, view_count, stop - start, 5 ) observer_valid = projected_valid & inside[:, None] observer_valid &= other_view_mask[reference_index][None, :, None, None] reference_valid = projected_valid[:, reference_index] & inside candidate_valid = reference_valid & observer_valid.any(dim=1) reference_feature = reference_features[ :, reference_index, start:stop ][:, None, :, None] reference_color = reference_rgb[:, reference_index, start:stop][ :, None, :, None ] feature_agreement = ( functional.cosine_similarity( reference_feature, sampled_features, dim=-1, eps=1e-6, ) + 1.0 ) * 0.5 rgb_error = (sampled_rgb - reference_color).square().mean(dim=-1) rgb_agreement = torch.exp(-rgb_error / evidence_temperature) pair_score = feature_agreement * rgb_agreement score = ( pair_score * observer_valid.to(pair_score.dtype) ).sum(dim=1) / (view_count - 1) if _telemetry_records is not None: _telemetry_records.append( ( score.detach().reshape(-1, score.shape[-1]).cpu(), candidate_valid.detach() .reshape(-1, candidate_valid.shape[-1]) .cpu(), observer_valid.detach() .sum(dim=1) .reshape(-1, candidate_valid.shape[-1]) .cpu(), ( None if appearance_points is None else appearance_points[:, start:stop].reshape(-1).cpu() ), ) ) peakness, contributes = _normalized_center_peakness( score, candidate_valid, temperature=surface_peak_temperature, ) peakness_sum[:, start:stop] += peakness reference_count[:, start:stop] += contributes.to(reference_count.dtype) peakness = peakness_sum / reference_count.clamp_min(1.0) peakness = torch.where(reference_count > 0, peakness, torch.zeros_like(peakness)) return _points_to_volume(peakness.unsqueeze(-1), resolution).detach() def _surface_peakness_with_telemetry( context_rgb: Tensor, cameras: Cameras, bounds: Tensor, resolution: tuple[int, int, int], *, evidence_temperature: float, surface_peak_temperature: float, chunk_size: int, appearance_confidence: Tensor | None = None, ) -> tuple[Tensor, dict[str, object]]: """Run the production cue while collecting context-only candidate statistics.""" records: list[tuple[Tensor, Tensor, Tensor, Tensor | None]] = [] peakness = _surface_peakness( context_rgb, cameras, bounds, resolution, evidence_temperature=evidence_temperature, surface_peak_temperature=surface_peak_temperature, chunk_size=chunk_size, _telemetry_records=records, _telemetry_appearance_confidence=appearance_confidence, ) if records: scores = torch.cat([record[0] for record in records], dim=0) valid = torch.cat([record[1] for record in records], dim=0) observer_count = torch.cat([record[2] for record in records], dim=0) profile_appearance = ( None if records[0][3] is None else torch.cat([record[3] for record in records if record[3] is not None], dim=0) ) else: scores = torch.empty(0, 5) valid = torch.empty(0, 5, dtype=torch.bool) observer_count = torch.empty(0, 5, dtype=torch.int64) profile_appearance = None report = _candidate_profile_telemetry( scores, valid, observer_count, offsets=torch.tensor([-2, -1, 0, 1, 2]), temperature=surface_peak_temperature, appearance_confidence=profile_appearance, observer_capacity=max(context_rgb.shape[1] - 1, 1), ) report["surface_peakness"] = _tensor_summary(peakness) return peakness, report def _normalized_center_peakness( scores: Tensor, valid: Tensor, *, temperature: float, center_index: int = 2, ) -> tuple[Tensor, Tensor]: """Normalize center probability so a flat valid candidate profile maps to zero.""" if scores.ndim < 1 or scores.shape != valid.shape: raise ValueError("scores and valid must have the same non-empty shape") if valid.dtype != torch.bool: raise ValueError("valid must be boolean") if not 0 <= center_index < scores.shape[-1]: raise ValueError("center_index must identify a candidate") if temperature <= 0: raise ValueError("temperature must be positive") if not torch.isfinite(scores).all(): raise ValueError("scores must be finite") candidate_count = valid.sum(dim=-1) contributes = valid[..., center_index] & (candidate_count >= 3) safe_valid = valid.clone() safe_valid[..., center_index] |= ~valid.any(dim=-1) logits = (scores.float() / temperature).masked_fill(~safe_valid, -torch.inf) center_probability = torch.softmax(logits, dim=-1)[..., center_index] count = candidate_count.to(center_probability.dtype) peakness = ((count * center_probability - 1.0) / (count - 1.0).clamp_min(1.0)).clamp( 0.0, 1.0 ) return torch.where(contributes, peakness, torch.zeros_like(peakness)), contributes def _positive_candidate_excess( scores: Tensor, valid: Tensor, *, temperature: float, min_candidates: int = 3, ) -> tuple[Tensor, Tensor]: """Return valid-only probability mass above the flat-profile baseline.""" if scores.ndim < 1 or scores.shape != valid.shape: raise ValueError("scores and valid must have the same non-empty shape") if valid.dtype != torch.bool: raise ValueError("valid must be boolean") if temperature <= 0: raise ValueError("temperature must be positive") if not 2 <= min_candidates <= scores.shape[-1]: raise ValueError("min_candidates must be between two and the candidate count") if not torch.isfinite(scores).all(): raise ValueError("scores must be finite") candidate_count = valid.sum(dim=-1) contributes = candidate_count >= min_candidates safe_valid = valid.clone() safe_valid[..., 0] |= ~valid.any(dim=-1) logits = (scores.float() / temperature).masked_fill(~safe_valid, -torch.inf) probabilities = torch.softmax(logits, dim=-1) count = candidate_count.to(probabilities.dtype) excess = ( (count[..., None] * probabilities - 1.0) / (count[..., None] - 1.0).clamp_min(1.0) ).clamp(0.0, 1.0) excess *= valid.to(excess.dtype) return torch.where(contributes[..., None], excess, torch.zeros_like(excess)), contributes def _candidate_profile_telemetry( scores: Tensor, valid: Tensor, observer_count: Tensor, *, offsets: Tensor, temperature: float, center_index: int = 2, appearance_confidence: Tensor | None = None, observer_capacity: int | None = None, ) -> dict[str, object]: """Summarize fixed depth-candidate competition without target labels.""" if scores.ndim != 2 or scores.shape != valid.shape: raise ValueError("scores and valid must have shape [N, K]") if observer_count.shape != scores.shape: raise ValueError("observer_count must align with scores") if valid.dtype != torch.bool: raise ValueError("valid must be boolean") if offsets.ndim != 1 or offsets.shape[0] != scores.shape[1]: raise ValueError("offsets must identify every candidate") if not 0 <= center_index < scores.shape[1]: raise ValueError("center_index must identify a candidate") if temperature <= 0: raise ValueError("temperature must be positive") if not torch.isfinite(scores).all(): raise ValueError("scores must be finite") if (observer_count < 0).any(): raise ValueError("observer_count cannot be negative") if observer_capacity is not None and observer_capacity < 1: raise ValueError("observer_capacity must be positive") if observer_capacity is not None and (observer_count > observer_capacity).any(): raise ValueError("observer_count cannot exceed observer_capacity") if appearance_confidence is not None and appearance_confidence.shape != scores.shape[:1]: raise ValueError("appearance_confidence must have shape [N]") scores = scores.detach().float().cpu() valid = valid.detach().cpu() observer_count = observer_count.detach().to(dtype=torch.int64, device="cpu") offsets = offsets.detach().float().cpu() if appearance_confidence is not None: appearance_confidence = appearance_confidence.detach().float().cpu() if not torch.isfinite(appearance_confidence).all(): raise ValueError("appearance_confidence must be finite") valid_count = valid.sum(dim=-1) contributes = valid[:, center_index] & (valid_count >= 3) profile_total = int(scores.shape[0]) contributing_total = int(contributes.sum().item()) report: dict[str, object] = { "profiles": { "total": profile_total, "contributing": contributing_total, "contributing_fraction": ( float(contributing_total / profile_total) if profile_total else 0.0 ), }, "valid_candidate_count_histogram": _integer_histogram(valid_count), "observer_count_histogram": _integer_histogram(observer_count[valid]), "competition": {}, "strata": {}, } if contributing_total == 0: return report report["competition"] = _competition_telemetry( scores[contributes], valid[contributes], offsets=offsets, temperature=temperature, center_index=center_index, ) if appearance_confidence is not None: positive = contributes & (appearance_confidence > 1e-6) top_threshold = float(torch.quantile(appearance_confidence, 0.9).item()) top_decile = contributes & (appearance_confidence >= top_threshold) report["strata"] = { "appearance_positive": _competition_stratum( scores, valid, positive, offsets=offsets, temperature=temperature, center_index=center_index, threshold=1e-6, ), "appearance_top_decile": _competition_stratum( scores, valid, top_decile, offsets=offsets, temperature=temperature, center_index=center_index, threshold=top_threshold, ), } if observer_capacity is not None: observer_mean_scores = ( scores * float(observer_capacity) / observer_count.to(dtype=scores.dtype).clamp_min(1.0) ) report["counterfactuals"] = { "valid_observer_mean": _candidate_profile_telemetry( observer_mean_scores, valid, observer_count, offsets=offsets, temperature=temperature, center_index=center_index, appearance_confidence=appearance_confidence, ), "min_two_observers": _candidate_profile_telemetry( observer_mean_scores, valid & (observer_count >= 2), observer_count, offsets=offsets, temperature=temperature, center_index=center_index, appearance_confidence=appearance_confidence, ), } return report def _competition_stratum( scores: Tensor, valid: Tensor, selected: Tensor, *, offsets: Tensor, temperature: float, center_index: int, threshold: float, ) -> dict[str, object]: profile_count = int(selected.sum().item()) return { "profiles": profile_count, "appearance_threshold": float(threshold), "competition": ( _competition_telemetry( scores[selected], valid[selected], offsets=offsets, temperature=temperature, center_index=center_index, ) if profile_count else {} ), } def _competition_telemetry( profile_scores: Tensor, profile_valid: Tensor, *, offsets: Tensor, temperature: float, center_index: int, ) -> dict[str, object]: profile_count = profile_valid.sum(dim=-1).float() masked_scores = profile_scores.masked_fill(~profile_valid, -torch.inf) logits = (profile_scores / temperature).masked_fill(~profile_valid, -torch.inf) probabilities = torch.softmax(logits, dim=-1) safe_probabilities = torch.where(profile_valid, probabilities, torch.zeros_like(probabilities)) maximum = masked_scores.max(dim=-1).values maximum_mask = profile_valid & torch.isclose( profile_scores, maximum[:, None], rtol=1e-6, atol=1e-7, ) maximum_count = maximum_mask.sum(dim=-1) unique_winner = maximum_count == 1 center_is_maximum = maximum_mask[:, center_index] unique_center = unique_winner & center_is_maximum center_fractional = maximum_mask[:, center_index].float() / maximum_count.float() chance_rate = 1.0 / profile_count other_valid = profile_valid.clone() other_valid[:, center_index] = False other_maximum = profile_scores.masked_fill(~other_valid, -torch.inf).max(dim=-1).values center_margin = profile_scores[:, center_index] - other_maximum minimum = profile_scores.masked_fill(~profile_valid, torch.inf).min(dim=-1).values score_range = maximum - minimum score_mean = ( torch.where(profile_valid, profile_scores, torch.zeros_like(profile_scores)).sum(dim=-1) / profile_count ) score_variance = ( torch.where( profile_valid, (profile_scores - score_mean[:, None]).square(), torch.zeros_like(profile_scores), ).sum(dim=-1) / profile_count ) entropy = -( safe_probabilities * safe_probabilities.clamp_min(torch.finfo(torch.float32).tiny).log() ).sum(dim=-1) / profile_count.log() sorted_probabilities = safe_probabilities.sort(dim=-1, descending=True).values positive_excess = ( (profile_count[:, None] * probabilities - 1.0) / (profile_count[:, None] - 1.0).clamp_min(1.0) ).clamp(0.0, 1.0) * profile_valid.to(probabilities.dtype) excess_per_profile = positive_excess.sum(dim=-1) noncenter_excess = positive_excess.clone() noncenter_excess[:, center_index] = 0.0 total_excess = float(excess_per_profile.sum().item()) winner_offsets = offsets[masked_scores.argmax(dim=-1)[unique_winner]] winner_histogram: dict[str, int] = {} for offset in offsets: count = int(torch.isclose(winner_offsets, offset).sum().item()) if count: winner_histogram[_number_key(float(offset.item()))] = count return { "center_argmax_inclusive_rate": float(center_is_maximum.float().mean().item()), "center_unique_argmax_rate": float(unique_center.float().mean().item()), "center_argmax_fractional_rate": float(center_fractional.mean().item()), "chance_argmax_rate": float(chance_rate.mean().item()), "center_argmax_excess_over_chance": float( (center_fractional - chance_rate).mean().item() ), "unique_winner_fraction": float(unique_winner.float().mean().item()), "winner_offset_histogram": winner_histogram, "noncenter_positive_excess_fraction": ( float(noncenter_excess.sum().item()) / total_excess if total_excess > 0 else 0.0 ), "zero_positive_excess_profile_count": int((excess_per_profile == 0).sum().item()), "positive_excess_total": total_excess, "center_margin": _tensor_summary(center_margin), "score_range": _tensor_summary(score_range), "score_std": _tensor_summary(score_variance.sqrt()), "softmax_entropy_normalized": _tensor_summary(entropy), "p_center": _tensor_summary(probabilities[:, center_index]), "p_max": _tensor_summary(sorted_probabilities[:, 0]), "p_second": _tensor_summary(sorted_probabilities[:, 1]), } def _integer_histogram(values: Tensor) -> dict[str, int]: values = values.detach().to(dtype=torch.int64, device="cpu").reshape(-1) if values.numel() == 0: return {} unique, counts = torch.unique(values, sorted=True, return_counts=True) return { str(int(value.item())): int(count.item()) for value, count in zip(unique, counts, strict=True) } def _number_key(value: float) -> str: return str(int(value)) if value.is_integer() else f"{value:g}" def _tensor_summary(values: Tensor) -> dict[str, float]: values = values.detach().float().cpu().reshape(-1) if values.numel() == 0: return {} quantiles = torch.quantile(values, torch.tensor([0.1, 0.5, 0.9])) return { "count": float(values.numel()), "mean": float(values.mean().item()), "std": float(values.std(unbiased=False).item()), "min": float(values.min().item()), "p10": float(quantiles[0].item()), "p50": float(quantiles[1].item()), "p90": float(quantiles[2].item()), "max": float(values.max().item()), } def _cross_view_evidence( features: Tensor, rgb: Tensor, valid: Tensor, *, temperature: float, ) -> tuple[Tensor, Tensor]: """Return correspondence confidence `[B,N,1]` and view provenance `[B,V,N,1]`.""" if features.ndim != 4 or rgb.ndim != 4 or valid.ndim != 3: raise ValueError("features, rgb, and valid must have [B,V,N,...] shapes") if features.shape[:3] != rgb.shape[:3] or features.shape[:3] != valid.shape: raise ValueError("features, rgb, and valid must align") if rgb.shape[-1] != 3: raise ValueError("rgb must have three channels") if temperature <= 0: raise ValueError("temperature must be positive") # Evidence audits the context; it must not become a trainable escape route for completion. features = features.detach() rgb = rgb.detach() mask = valid.unsqueeze(-1).to(features.dtype) count = mask.sum(dim=1).clamp_min(1.0) normalized = functional.normalize(features, dim=-1, eps=1e-6) * mask feature_sum = normalized.sum(dim=1) pair_dot_sum = feature_sum.square().sum(dim=-1, keepdim=True) - mask.sum(dim=1) raw_count = mask.sum(dim=1) pair_count = (raw_count * (raw_count - 1.0)).clamp_min(1.0) pair_correlation = (pair_dot_sum / pair_count).clamp(-1.0, 1.0) feature_agreement = (pair_correlation + 1.0) * 0.5 rgb_mean = (rgb * mask).sum(dim=1) / count rgb_variance = ((rgb - rgb_mean[:, None]).square() * mask).sum(dim=1) / count rgb_agreement = torch.exp(-rgb_variance.mean(dim=-1, keepdim=True) / temperature) view_count = features.shape[1] if view_count == 1: pair_support = torch.zeros_like(raw_count) else: pair_support = ((raw_count - 1.0) / (view_count - 1.0)).clamp(0.0, 1.0) confidence = (pair_support * feature_agreement * rgb_agreement).clamp(0.0, 1.0) consensus = functional.normalize(feature_sum, dim=-1, eps=1e-6) feature_score = (normalized * consensus[:, None]).sum(dim=-1, keepdim=True) rgb_score = -(rgb - rgb_mean[:, None]).square().mean(dim=-1, keepdim=True) / temperature score = feature_score / temperature + rgb_score score = score.masked_fill(~valid.unsqueeze(-1), -1e4) provenance = torch.softmax(score, dim=1) * valid.unsqueeze(-1).to(score.dtype) has_valid = valid.any(dim=1, keepdim=True).unsqueeze(-1) provenance = provenance * has_valid.to(provenance.dtype) provenance = provenance / provenance.sum(dim=1, keepdim=True).clamp_min(1e-8) return confidence, provenance def _project_voxel_grid( bounds: Tensor, cameras: Cameras, resolution: tuple[int, int, int], image_height: int, image_width: int, ) -> tuple[Tensor, Tensor, Tensor, Tensor]: centers = make_voxel_centers(bounds, resolution) point_count = math.prod(resolution) flat_centers = centers.reshape(bounds.shape[0], point_count, 3) points_per_view = flat_centers[:, None].expand( bounds.shape[0], cameras.leading_shape[1], point_count, 3 ) with torch.autocast(device_type=bounds.device.type, enabled=False): pixels, camera_depth, valid = project_world( points_per_view.float(), cameras.to(dtype=torch.float32) ) normalized_grid = _pixels_to_grid(pixels, image_height, image_width) return points_per_view, normalized_grid, camera_depth, valid def _base_color_from_provenance(sampled_rgb: Tensor, provenance: Tensor) -> Tensor: base_color = (sampled_rgb * provenance).sum(dim=1) has_provenance = provenance.sum(dim=1) > 0.0 return torch.where( has_provenance.expand_as(base_color), base_color, torch.full_like(base_color, 0.5), ) def _sample_images(images: Tensor, normalized_grid: Tensor) -> Tensor: batch_size, view_count, channels, _, _ = images.shape point_count = normalized_grid.shape[-2] flattened_images = images.reshape(batch_size * view_count, channels, *images.shape[-2:]) grid = normalized_grid.reshape(batch_size * view_count, point_count, 1, 2) sampled = functional.grid_sample( flattened_images, grid, mode="bilinear", padding_mode="zeros", align_corners=True ) return sampled.reshape(batch_size, view_count, channels, point_count).permute(0, 1, 3, 2) def _fixed_rgb_pyramid(images: Tensor, normalized_grid: Tensor) -> Tensor: batch_size, view_count, channels, height, width = images.shape flattened = images.reshape(batch_size * view_count, channels, height, width) levels = [flattened] for scale in (2, 4): if min(height, width) >= scale: levels.append(functional.avg_pool2d(flattened, kernel_size=scale, stride=scale)) sampled = [ _sample_images( level.reshape(batch_size, view_count, channels, *level.shape[-2:]), normalized_grid, ) for level in levels ] return torch.cat(sampled, dim=-1) def _pixels_to_grid(pixels: Tensor, height: int, width: int) -> Tensor: x = pixels[..., 0] / max(width - 1, 1) * 2.0 - 1.0 y = pixels[..., 1] / max(height - 1, 1) * 2.0 - 1.0 return torch.stack((x, y), dim=-1) def _points_to_volume(values: Tensor, resolution: tuple[int, int, int]) -> Tensor: return values.reshape(values.shape[0], *resolution, values.shape[-1]).permute(0, 4, 1, 2, 3) def _normalize(values: Tensor) -> Tensor: return values / torch.linalg.vector_norm(values, dim=-1, keepdim=True).clamp_min(1e-8) def _group_count(channels: int) -> int: for groups in (8, 4, 2): if channels % groups == 0: return groups return 1