fastkmeans 0.2.0__tar.gz → 0.3.0__tar.gz

{fastkmeans-0.2.0 → fastkmeans-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fastkmeans
-Version: 0.2.0
+Version: 0.3.0
 Summary: Add your description here
 Author-email: Ben Clavié <bc@answer.ai>, Benjamin Warner <bw@answer.ai>
 Maintainer-email: Ben Clavié <bc@answer.ai>, Benjamin Warner <bw@answer.ai>

{fastkmeans-0.2.0 → fastkmeans-0.3.0}/fastkmeans/__init__.py

@@ -1,4 +1,4 @@
 from .kmeans import FastKMeans
 
 __all__ = ["FastKMeans"]
-__version__ = "0.2.0"
+__version__ = "0.3.0"

{fastkmeans-0.2.0 → fastkmeans-0.3.0}/fastkmeans/kmeans.py

@@ -3,8 +3,6 @@ import time
 import torch
 import numpy as np
 
-from fastkmeans.triton_kernels import chunked_kmeans_kernel
-
 def _get_device(preset: str | int | torch.device | None = None):
     if isinstance(preset, torch.device):
         return preset
@@ -53,6 +51,9 @@ def _kmeans_torch_double_chunked(
         Where n_samples_used can be smaller than the original if subsampling occurred.
     """
 
+    if use_triton:
+        from fastkmeans.triton_kernels import chunked_kmeans_kernel
+
     if dtype is None:
         dtype = torch.float16 if device.type in ['cuda', 'xpu'] else torch.float32
 
@@ -274,6 +275,8 @@ class FastKMeans:
         -------
         labels : np.ndarray of shape (n_samples,), int64
         """
+        if self.use_triton:
+            from fastkmeans.triton_kernels import chunked_kmeans_kernel
         if self.centroids is None:
             raise RuntimeError("Must call train() or fit() before predict().")
 

{fastkmeans-0.2.0 → fastkmeans-0.3.0}/fastkmeans.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fastkmeans
-Version: 0.2.0
+Version: 0.3.0
 Summary: Add your description here
 Author-email: Ben Clavié <bc@answer.ai>, Benjamin Warner <bw@answer.ai>
 Maintainer-email: Ben Clavié <bc@answer.ai>, Benjamin Warner <bw@answer.ai>

{fastkmeans-0.2.0 → fastkmeans-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fastkmeans"
-version = "0.2.0"
+version = "0.3.0"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.8"