Benchmark np.linalg.multi_dot on dense matrices for scikit-learn/scikit-learn#19571

multi_dot_logreg.py

import numpy as np

from .common import Benchmark


class MultiDotLogReg(Benchmark):
    """ Benchmark np.linalg.multi_dot on dense matrices for #19571
    """

    param_names = [
        'memory_layout',
        'n_samples',
        'n_features',
        's_n_cols_prop',
        'mult_strategy'
    ]

    params = (
        ['C', 'F'],
        [100, 1_000, 10_000, 100_000],
        [100, 1_000],
        [0.1, 0.5, 1, 2],
        ['use_multi_dot', 'current_order', 'alternative_order'],
    )

    def setup(self, memory_layout, n_samples, n_features,
              s_n_cols_prop, mult_strategy):
        self.X = np.array(np.random.rand(n_samples, n_features),
                          order=memory_layout)
        self.dX = np.array(np.random.rand(n_samples, n_features),
                           order=memory_layout)
        # s[:n_features]
        s_n_cols = int(n_features * s_n_cols_prop)
        self.s_trimmed = np.array(np.random.rand(n_features, s_n_cols),
                                  order=memory_layout)

        self.mult_strategy = mult_strategy

    def time_hessian_grad_prod(self, memory_layout,n_samples, n_features,
                               s_n_cols_prop, mult_strategy):
        if self.mult_strategy == 'use_multi_dot':
            # proposed modification
            np.linalg.multi_dot([self.X.T, self.dX, self.s_trimmed])
        elif self.mult_strategy == 'current_order':
            self.X.T.dot(self.dX.dot(self.s_trimmed))
        else: # alternative_order
            self.X.T.dot(self.dX).dot(self.s_trimmed)