Added shift.

Author: KateBlueSky

configs/spmd/kmeans_wide_weak.json

@@ -25,7 +25,7 @@
             "SETS": [
                 "synthetic data",
                 "sklearnex spmd implementation",
-                "large scale 2k parameters sample shift",
+                "large scale 2k parameters",
                 "spmd kmeans parameters"
             ]
         }

sklbench/datasets/transformer.py

@@ -23,7 +23,7 @@
 
 from ..utils.bench_case import get_bench_case_value
 from ..utils.logger import logger
-
+from mpi4py import MPI
 
 def convert_data(data, dformat: str, order: str, dtype: str, device: str = None):
     if isinstance(data, csr_matrix) and dformat != "csr_matrix":
@@ -113,8 +113,36 @@ def split_and_transform_data(bench_case, data, data_description):
         "KNeighbors" in get_bench_case_value(bench_case, "algorithm:estimator", "")
         and int(get_bench_case_value(bench_case, "bench:mpi_params:n", 1)) > 1
     )
-    if distributed_split == "rank_based" or knn_split_train:
-        from mpi4py import MPI
+
+    if distributed_split == "sample_shift":
+       comm = MPI.COMM_WORLD
+       rank = comm.Get_rank()
+       size = comm.Get_size()  
+
+       n_train = len(x_train)
+       n_test = len(x_test)
+
+       train_start = 0
+       train_end = n_train
+       test_start = 0
+       test_end = n_test 
+
+       adjust_number = (math.sqrt(rank) * 0.003) + 1
+
+       if "y" in data:
+            x_train, y_train = (
+                x_train[train_start:train_end] *  adjust_number,
+                y_train[train_start:train_end],
+            )
+            
+            x_test, y_test = x_test[test_start:test_end] * adjust_number, y_test[test_start:test_end]
+       else:
+            x_train = x_train[train_start:train_end]
+        
+            x_test = x_test[test_start:test_end] * adjust_number
+
+    elif distributed_split == "rank_based" or knn_split_train:
+        
 
         comm = MPI.COMM_WORLD
         rank = comm.Get_rank()
@@ -127,6 +155,7 @@ def split_and_transform_data(bench_case, data, data_description):
         train_end = (1 + rank) * n_train // size
         test_start = rank * n_test // size
         test_end = (1 + rank) * n_test // size
+        x_train_rank = x_train[train_start:train_end] 
 
         if "y" in data:
             x_train, y_train = (
@@ -138,7 +167,7 @@ def split_and_transform_data(bench_case, data, data_description):
         else:
             x_train = x_train[train_start:train_end]
             if distributed_split == "rank_based":
-                x_test = x_test[test_start:test_end]
+                x_test = x_test[test_start:test_end] * adjust_number
 
     device = get_bench_case_value(bench_case, "algorithm:device", None)
     common_data_format = get_bench_case_value(bench_case, "data:format", "pandas")