def forward(self, geometry_feat, rgb_feat, ray_diff, mask):
        '''
        :param geometry_feat: geometry features indicates sdf  [n_rays, n_samples, n_feat]
        :param rgb_feat: rgbs and image features [n_views, n_rays, n_samples, n_feat]
        :param ray_diff: ray direction difference [n_views, n_rays, n_samples, 4], first 3 channels are directions,
        last channel is inner product
        :param mask: mask for whether each projection is valid or not. [n_views, n_rays, n_samples]
        :return: rgb and density output, [n_rays, n_samples, 4]
        '''

        rgb_feat = rgb_feat.permute(1, 2, 0, 3).contiguous()
        ray_diff = ray_diff.permute(1, 2, 0, 3).contiguous()
        mask = mask[:, :, :, None].permute(1, 2, 0, 3).contiguous()
        num_views = rgb_feat.shape[2]
        geometry_feat = geometry_feat[:, :, None, :].repeat(1, 1, num_views, 1)

        direction_feat = self.ray_dir_fc(ray_diff)
        rgb_in = rgb_feat[..., :3]
        rgb_feat = rgb_feat + direction_feat

        if self.anti_alias_pooling:
            _, dot_prod = torch.split(ray_diff, [3, 1], dim=-1)
            exp_dot_prod = torch.exp(torch.abs(self.s) * (dot_prod - 1))
            weight = (exp_dot_prod - torch.min(exp_dot_prod, dim=2, keepdim=True)[0]) * mask
            weight = weight / (torch.sum(weight, dim=2, keepdim=True) + 1e-8)
        else:
            weight = mask / (torch.sum(mask, dim=2, keepdim=True) + 1e-8)

        # compute mean and variance across different views for each point
        mean, var = fused_mean_variance(rgb_feat, weight)  # [n_rays, n_samples, 1, n_feat]
        globalfeat = torch.cat([mean, var], dim=-1)  # [n_rays, n_samples, 1, 2*n_feat]

        x = torch.cat([geometry_feat, globalfeat.expand(-1, -1, num_views, -1), rgb_feat],
                      dim=-1)  # [n_rays, n_samples, n_views, 3*n_feat+n_geo_feat]
        x = self.base_fc(x)

        x_vis = self.vis_fc(x * weight)
        x_res, vis = torch.split(x_vis, [x_vis.shape[-1] - 1, 1], dim=-1)
        vis = torch.sigmoid(vis) * mask
        x = x + x_res
        vis = self.vis_fc2(x * vis) * mask

        # rgb computation
        x = torch.cat([x, vis, ray_diff], dim=-1)
        x = self.rgb_fc(x)
        x = x.masked_fill(mask == 0, -1e9)
        blending_weights_valid = F.softmax(x, dim=2)  # color blending
        rgb_out = torch.sum(rgb_in * blending_weights_valid, dim=2)

        mask = mask.detach().to(rgb_out.dtype)  # [n_rays, n_samples, n_views, 1]
        mask = torch.sum(mask, dim=2, keepdim=False)
        mask = mask >= 2  # more than 2 views see the point
        mask = torch.sum(mask.to(rgb_out.dtype), dim=1, keepdim=False)
        valid_mask = mask > 8  # valid rays, more than 8 valid samples
        return rgb_out, valid_mask  # (N_rays, n_samples, 3), (N_rays, 1)

        :param geometry_feat: geometry features indicates sdf  [n_rays, n_samples, n_feat]
        :param rgb_feat: rgbs and image features [n_views, n_rays, n_samples, n_feat]
        :param ray_diff: ray direction difference [n_views, n_rays, n_samples, 4], first 3 channels are directions,
        last channel is inner product
        :param mask: mask for whether each projection is valid or not. [n_views, n_rays, n_samples]
        :return: rgb and density output, [n_rays, n_samples, 4]
        
mv-diffusion/One-2-3-45/reconstruction/models/rendering_network.py