yobx.sklearn.cluster.mean_shift

yobx.sklearn.cluster.mean_shift.sklearn_mean_shift(g: GraphBuilderExtendedProtocol, sts: Dict, outputs: List[str], estimator: MeanShift, X: str, name: str = 'mean_shift') → str | Tuple[str, str]

Converts a sklearn.cluster.MeanShift into ONNX.

The converter produces two outputs: the predicted cluster labels (equivalent to predict()) and the Euclidean distances from each sample to every cluster centre (equivalent to the distance matrix used internally by predict).

The predict method of MeanShift assigns each sample to the nearest cluster centre using Euclidean distance, which is reproduced here via the identity:

||x - c||² = ||x||² - 2·x·cᵀ + ||c||²

Two computation paths are used depending on the available opsets:

With com.microsoft opset (CDist path):

X (N,F)  centers (K,F)
      │       │
 com.microsoft.CDist(metric="euclidean") ──► distances (N,K)
                                                  │
                                       ArgMin(axis=1) ──► labels (N,)

Without com.microsoft opset (standard ONNX path):

X (N,F)
  │
  ├──Mul──ReduceSum(axis=1, keepdims=1)──────────────────────────────► x_sq (N,1)
  │                                                                         │
  └──MatMul(centersᵀ)────────────────────────────────────────────────► cross (N,K)
                                                                            │
c_sq (1,K) ─────────────────────── Add(x_sq) ─── Sub(Mul(2,cross)) ──► sq_dists (N,K)
                                                                            │
                                           Sqrt ──────────────────────► distances (N,K)
                                                                            │
                                       ArgMin(axis=1) ─────────────────► labels (N,)

Parameters:

• g – the graph builder to add nodes to

• sts – shapes defined by scikit-learn

• estimator – a fitted MeanShift

• outputs – desired output names; outputs[0] receives the cluster labels and outputs[1] (if present) receives the distances matrix

• X – input tensor name

• name – prefix names for the added nodes

Returns:

tuple (labels, distances) when two outputs are requested, otherwise just labels