fastkmeans 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

fastkmeans/__init__.py

@@ -1,4 +1,4 @@
 from .kmeans import FastKMeans
 
 __all__ = ["FastKMeans"]
-__version__ = "0.2.0"
+__version__ = "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.dist-info → fastkmeans-0.3.0.dist-info}/METADATA

@@ -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.3.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+fastkmeans/__init__.py,sha256=EKvVznV3rTwe5Yz_tqlF_bJMZsh-5vLSjVvMtCx3Xhk,79
+fastkmeans/kmeans.py,sha256=_FbU_avJPxIEfikhPMPgBibj9ckKej21iWJCm-YNEUM,13778
+fastkmeans/triton_kernels.py,sha256=iN8khhoaQGJ08LQy5iz4VGEXRCtvFKDfCgyGzwVqjgw,3698
+fastkmeans-0.3.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+fastkmeans-0.3.0.dist-info/METADATA,sha256=KOyw3KsNPCdF_6spJl2SByB8mVtZ3I6GE1H8N2KMUpM,6791
+fastkmeans-0.3.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+fastkmeans-0.3.0.dist-info/top_level.txt,sha256=B3Zd2-kEAH_hN0hFUWgo5lO-TH7ppVol_WQ5ZT1H0js,11
+fastkmeans-0.3.0.dist-info/RECORD,,

{fastkmeans-0.2.0.dist-info → fastkmeans-0.3.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (78.1.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

fastkmeans-0.2.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-fastkmeans/__init__.py,sha256=mPLGfhksBBfB1dTkPBPrPg5E1qixixl_Qritc6A10AI,79
-fastkmeans/kmeans.py,sha256=RLDeaCIbKpCJBzV2XO3JPgLwjrhJ6vtA997Gyzm_GyA,13651
-fastkmeans/triton_kernels.py,sha256=iN8khhoaQGJ08LQy5iz4VGEXRCtvFKDfCgyGzwVqjgw,3698
-fastkmeans-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-fastkmeans-0.2.0.dist-info/METADATA,sha256=2hsTZr0t2_CNhCyrECIBk2S4ZB7ytF0YaDqgcgHDvDc,6791
-fastkmeans-0.2.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-fastkmeans-0.2.0.dist-info/top_level.txt,sha256=B3Zd2-kEAH_hN0hFUWgo5lO-TH7ppVol_WQ5ZT1H0js,11
-fastkmeans-0.2.0.dist-info/RECORD,,