Repository: scikit-learn/scikit-learn
File: sklearn/cluster/_kmeans.py (Python)

Code before:
                best_inertia = inertia

        centers = init_centers
        centers_new = np.empty_like(centers)

        # Initialize counts
        self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

        # Attributes to monitor the convergence
        self._ewa_inertia = None
        self._ewa_inertia_min = None
        self._no_improvement = 0

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

        n_steps = (self.max_iter * n_samples) // self._batch_size

        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            # Perform the iterative optimization until convergence
            for i in range(n_steps):
                # Sample a minibatch from the full dataset
                minibatch_indices = random_state.choice(
                    n_samples,
                    self._batch_size,
                    p=sample_weight / np.sum(sample_weight),
                    replace=True,
                )

                # Perform the actual update step on the minibatch data
                # Note, I am not sure how sample weights are used here
                # So left it in, it seems like the weight sums are updated using
                # sample weights so need some help here to understand the
                # _minibatch_update_dense/sparse code
                batch_inertia = _mini_batch_step(
                    X=X[minibatch_indices],
                    sample_weight=sample_weight[minibatch_indices],
                    centers=centers,
                    centers_new=centers_new,
                    weight_sums=self._counts,
                    random_state=random_state,
                    random_reassign=self._random_reassign(),
                    reassignment_ratio=self.reassignment_ratio,
                    verbose=self.verbose,
                    n_threads=self._n_threads,
                )

                if self._tol > 0.0:
                    centers_squared_diff = np.sum((centers_new - centers) ** 2)
                else:
                    centers_squared_diff = 0

Review comment (performance):
For performance reasons, we should compute `sw / sw.sum()` outside of the loop since it's always the same.

Code after:
            cluster_centers = self._init_centroids(
                X,
                x_squared_norms=x_squared_norms,
                init=init,
                random_state=random_state,
                init_size=self._init_size,
                sample_weight=sample_weight,
            )

            # Compute inertia on a validation set.
            _, inertia = _labels_inertia_threadpool_limit(
                X_valid,
                sample_weight_valid,
                cluster_centers,
                n_threads=self._n_threads,
            )

            if self.verbose:
                print(f"Inertia for init {init_idx + 1}/{self._n_init}: {inertia}")
            if best_inertia is None or inertia < best_inertia:
                init_centers = cluster_centers
                best_inertia = inertia

        centers = init_centers
        centers_new = np.empty_like(centers)

        # Initialize counts
        self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

        # Attributes to monitor the convergence
        self._ewa_inertia = None
        self._ewa_inertia_min = None
        self._no_improvement = 0

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

        sum_of_weights = np.sum(sample_weight)

        n_steps = (self.max_iter * n_samples) // self._batch_size
        normalized_sample_weight = sample_weight / sum_of_weights
        unit_sample_weight = np.ones_like(sample_weight, shape=(self._batch_size,))

        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            # Perform the iterative optimization until convergence
            for i in range(n_steps):
                # Sample a minibatch from the full dataset
                minibatch_indices = random_state.choice(
                    n_samples,
                    self._batch_size,
                    p=normalized_sample_weight,