def predict(
        self,
        features: torch.Tensor,
        group: str | None = None,
        n_neighbors: int = 1,
        normalize: bool = True,
    ) -> torch.Tensor:
        """Predict the (normalized) density for a set of features.

        Args:
            features (torch.Tensor): Input features of shape ``(N, D)`` that will be
                compared to the density model.
            group (str | None, optional): Group (video id) from which the features
                originate. If passed, all features of the same group in the memory
                bank will be excluded from the density estimation.
                Defaults to ``None``.
            n_neighbors (int, optional): Number of neighbors used in the KNN search.
                Defaults to ``1``.
            normalize (bool, optional): Flag indicating if the density should be
                normalized to min-max stats of the feature bank.
                Defaults to ``True``.

        Returns:
            torch.Tensor: Mean (normalized) distances of input feature vectors to k
                nearest neighbors in feature bank.

        Example:
            >>> estimator = GroupedKNNEstimator(n_neighbors=5)
            >>> features = torch.randn(32, 512)  # (N, D)
            >>> estimator.update(features, group="video_1")
            >>> estimator.fit()
            >>> scores = estimator.predict(features, group="video_1")
            >>> scores.shape
            torch.Size([32])
        """
        n_neighbors = n_neighbors or self.n_neighbors

        if group:
            group_idx = self.group_names.index(group)
            mem_bank = self.memory_bank[self.group_index != group_idx]
        else:
            mem_bank = self.memory_bank

        distances = self._nearest_neighbors(mem_bank, features, n_neighbors=n_neighbors)

        if normalize:
            distances = self._normalize(distances)

        return distances.mean(axis=1)
Predict the (normalized) density for a set of features.

        Args:
            features (torch.Tensor): Input features of shape ``(N, D)`` that will be
                compared to the density model.
            group (str | None, optional): Group (video id) from which the features
                originate. If passed, all features of the same group in the memory
                bank will be excluded from the density estimation.
                Defaults to ``None``.
            n_neighbors (int, optional): Number of neighbors used in the KNN search.
                Defaults to ``1``.
            normalize (bool, optional): Flag indicating if the density should be
                normalized to min-max stats of the feature bank.
                Defaults to ``True``.

        Returns:
            torch.Tensor: Mean (normalized) distances of input feature vectors to k
                nearest neighbors in feature bank.

        Example:
            >>> estimator = GroupedKNNEstimator(n_neighbors=5)
            >>> features = torch.randn(32, 512)  # (N, D)
            >>> estimator.update(features, group="video_1")
            >>> estimator.fit()
            >>> scores = estimator.predict(features, group="video_1")
            >>> scores.shape
            torch.Size([32])
        
src/anomalib/models/video/ai_vad/density.py