code: def stratify_by_features(features, n_strata, **kwargs):
      
      n_items = features.shape[0]
      km = KMeans(n_clusters=n_strata, **kwargs)
      allocations = km.fit_predict(X=features)
      return Strata(allocations)
docs: Stratify by clustering the items in feature space
  
      Parameters
      ----------
      features : array-like, shape=(n_items,n_features)
          feature matrix for the pool, where rows correspond to items and columns
          correspond to features.
  
      n_strata : int
          number of strata to create.
  
      **kwargs :
          passed to sklearn.cluster.KMeans
  
      Returns
      -------
      Strata instance