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

Code before:
                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

        n_steps = (self.max_iter * n_samples) // self._batch_size
        sample_weight = sample_weight / np.sum(sample_weight)
        sample_weight_minibatch_step = np.ones(self._batch_size).astype(X.dtype)
        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,
                    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_step,
                    centers=centers,
                    centers_new=centers_new,
                    weight_sums=self._counts,
                    random_state=random_state,
                    random_reassign=self._random_reassign(),
                    reassignment_ratio=self.reassignment_ratio,

Review comment (performance):
We need to keep an unnormalized version of the sample weights because they're used afterwards to compute the final inertia which scales with the sample weights.

```suggestion
        normalized_sample_weight = sample_weight / np.sum(sample_weight)
        unit_sample_weight = np.ones_like(sample_weights, shape=(self._batch_size,))
```

Code after:
                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,
                    replace=True,
                )
                # Perform the actual update step on the minibatch data
                # Note: since the sampling of the minibatch is sample_weight aware,