<jupyter_start><jupyter_text>Strassen matrix multiplication algorithm<jupyter_code>import numpy as np
from numba import njit<jupyter_output><empty_output><jupyter_text>Test data<jupyter_code>n = 1024
A = np.random.randn(n, n)
B = np.random.randn(n, n)<jupyter_output><empty_output><jupyter_text>Benchmark<jupyter_code>%timeit A @ B<jupyter_output>10 loops, best of 3: 21.3 ms per loop<jupyter_text>Naive IKJ multiplication<jupyter_code>@njit
def matrix_multiply(A, B):
   
    n = A.shape[0]
    C = np.zeros_like(A)

    for i in range(n):
        for k in range(n):       
            for j in range(n):
                C[i][j] += A[i][k] * B[k][j]
    return C
%timeit matrix_multiply(A, B)
np.allclose(A @ B, matrix_multiply(A, B))<jupyter_output><empty_output><jupyter_text>Naive strassen<jupyter_code>@njit
def matrix_merge(A, B, C, D):

    n = A.shape[0]
    result = np.empty(shape=(2*n, 2*n))

    for i in range(n):
        for j in range(n):
            result[i][j] = A[i][j]
            result[i][j + n] = B[i][j]
            result[i + n][j] = C[i][j]
            result[i + n][j + n] = D[i][j]

    return result


@njit(parallel=True)
def matrix_split(M):
  
    n = M.shape[0] // 2
 
    A = np.empty(shape=(n, n))
    B = np.empty_like(A)
    C = np.empty_like(A)
    D = np.empty_like(A)
    
    for i in range(n):    
        for j in range(n):
            A[i][j] = M[i][j]
            B[i][j] = M[i][j + n]
            C[i][j] = M[i + n][j]
            D[i][j] = M[i + n][j + n]

    return A, B, C, D

@njit
def matrix_add(A, B):

    n = A.shape[0]
    
    C = np.empty(shape=(n, n))
    for i in range(n):
        for j in range(n):
            C[i][j] = A[i][j] + B[i][j]
            
    return C


@njit
def matrix_subtract(A, B):

    n = A.shape[0]
    
    C = np.empty(shape=(n, n))
    for i in range(n):
        for j in range(n):
            C[i][j] = A[i][j] - B[i][j]
            
    return C


@njit
def strassen_naive(A, B):
    
    m = A.shape[0]

    A11, A12, A21, A22 = matrix_split(A)
    B11, B12, B21, B22 = matrix_split(B)
    S1 = matrix_subtract(B12, B22)
    S2 = matrix_add(A11, A12)
    S3 = matrix_add(A21, A22)
    S4 = matrix_subtract(B21, B11)
    S5 = matrix_add(A11, A22)
    S6 = matrix_add(B11,B22)
    S7 = matrix_subtract(A12, A22)
    S8 = matrix_add(B21, B22)
    S9 = matrix_subtract(A11, A21)
    S10 = matrix_add(B11, B12)

    cutoff = 256

    if m > cutoff:
        P1 = strassen_naive(A11, S1)
        P2 = strassen_naive(S2, B22)        
        P3 = strassen_naive(S3, B11)        
        P4 = strassen_naive(A22, S4)        
        P5 = strassen_naive(S5, S6)        
        P6 = strassen_naive(S7, S8)        
        P7 = strassen_naive(S9, S10)        
    else:
        P1 = matrix_multiply(A11, S1)
        P2 = matrix_multiply(S2, B22)
        P3 = matrix_multiply(S3, B11)
        P4 = matrix_multiply(A22, S4)
        P5 = matrix_multiply(S5, S6)
        P6 = matrix_multiply(S7, S8)
        P7 = matrix_multiply(S9, S10)

    C11 = matrix_add(matrix_add(P5, P6), matrix_subtract(P4, P2))
    C12 = matrix_add(P1, P2)
    C21 = matrix_add(P3, P4)
    C22 = matrix_subtract(matrix_add(P5, P1), matrix_add(P3, P7))

    return matrix_merge(C11, C12, C21, C22)

%timeit strassen_naive(A, B)
np.allclose(A @ B, strassen_naive(A, B))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14814814814814814, 'fraction_numerical': 0.0504778972520908, 'mean_word_length': 2.3191278493557976, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '106709', 'license': 'MIT', 'n_tokens_mistral': 1543, 'n_tokens_neox': 1458, 'n_words': 300, 'path': 'performance/2017_07_strassens_algorithm/notebooks/Strassen Algorithm.ipynb', 'repo_name': 'fastats/learning', 'size': '3348'}