def _local_transport(self, var_k: np.ndarray, x: np.ndarray, kernel_size: int) -> np.ndarray:
        """
        Compute local transport.

        :param var_k: K parameter in Algorithm 2 of the paper ``Wasserstein Adversarial Examples via Projected
            Sinkhorn Iterations``.
        :param x: An array to apply local transport.
        :param kernel_size: Kernel size for computing the cost matrix.
        :return: Local transport result.
        """
        # Compute number of channels
        num_channels = x.shape[1 if self.estimator.channels_first else 3]

        # Expand channels
        var_k = np.repeat(var_k, num_channels, axis=1)

        # Swap channels to prepare for local transport computation
        if not self.estimator.channels_first:
            x = np.swapaxes(x, 1, 3)

        # Compute local transport
        unfold_x = self._unfold(x=x, kernel_size=kernel_size, padding=kernel_size // 2)
        unfold_x = unfold_x.swapaxes(-1, -2)
        unfold_x = unfold_x.reshape(*unfold_x.shape[:-1], num_channels, kernel_size**2)
        unfold_x = unfold_x.swapaxes(-2, -3)

        tmp_k = var_k.reshape(var_k.shape[0], num_channels, -1)
        tmp_k = np.expand_dims(tmp_k, -1)

        result = np.matmul(unfold_x, tmp_k)
        result = np.squeeze(result, -1)
        result = result.reshape(*result.shape[:-1], x.shape[-2], x.shape[-1])

        # Swap channels for final result
        if not self.estimator.channels_first:
            result = np.swapaxes(result, 1, 3)

        return result

        Compute local transport.

        :param var_k: K parameter in Algorithm 2 of the paper ``Wasserstein Adversarial Examples via Projected
            Sinkhorn Iterations``.
        :param x: An array to apply local transport.
        :param kernel_size: Kernel size for computing the cost matrix.
        :return: Local transport result.
        
art/attacks/evasion/wasserstein.py