Custom Matrix Product in NumPy with User-Defined Element-Wise Operations

I am looking for a way to perform a custom matrix product in NumPy where I can replace the element-wise multiplication and addition operations with self-defined functions, such as lambda expressions. Specifically, I want to maintain the matrix multiplication signature (n,k),(k,m) -> (n,m), but use custom operators for the scalar multiplication and addition within the matrix product computation.

Example:

Given two matrices A and B:

A = [[a_11, a_12],
     [a_21, a_22]]

B = [[b_11, b_12],
     [b_21, b_22]]

The standard matrix product C = AB is calculated as:

C_ij = sum(A_ix * B_xj for x in range(k))

I want to replace the multiplication (*) and addition (sum) operations with custom functions, e.g., custom_multiply and custom_add, such that:

C_ij = custom_add(custom_multiply(A_i1, B_1j), custom_multiply(A_i2, B_2j), ..., custom_multiply(A_ik, B_kj))

Custom Operators:

For instance, I might want to define:

custom_multiply = lambda x, y: (x&0xff) * y  # Custom multiplication function
custom_add = lambda x, y: max(x, y)          # Custom addition function

Use Case:

I believe it is useful to have a custom matrix product for various applications. For example, in networking, the computation of routes can follow the general process of matrix multiplication when the adjacency matrix is given. However, the element-wise operations must be customized to conform to specific restrictions required by certain routing protocols.

Question:

I do not want to use plain loops to fulfill matrix product. My custom element-wise operations can be decomposed into a series of basic arithmetic operations that are supported by NumPy and can be extended/broadcast to matrix-wise. I'd like to know:

• Is there a built-in NumPy function to support custom matrix products?
• If not, is there a general way to break down the matrix product process into a series of matrix-wise operations in an efficient way?

Any guidance or examples would be greatly appreciated.

I am not aware of a readily available way to provide your own reduction and element-wise multiplication to the matrix product in Numpy. However, you could write your own function that achieves the same:

def my_matmul(a, b, my_sum, my_multiply):
    tmp = my_multiply(a[..., None], b[None, ...])  # n×k(×1) ∘ (1×)k×m → n×k×m
    return my_sum(tmp)  # n×k×m → n×m

Here, my_sum is the custom reduction function in place of the matrix product's sum, and my_multiply is the custom element-wise function in place of the matrix product's multiplication.

In the given form, my_sum and my_multiply have to fulfill the following criteria:

• my_sum always has to reduce along axis 1,
• my_multiply has to adhere to the standard Numpy broadcasting rules.

For example, we could use:

• my_sum = lambda array: np.max(array, axis=1) (Maximum instead of summation)
• my_multiply = lambda a, b: (a + 1) / (b - 1) (add 1 to elements in a, subtract 1 from elements in b, then divide instead of multiply)

Here is a full example, including test cases:

import numpy as np
rand = np.random.default_rng(42)

a, b = rand.normal(size=(3, 4)), rand.normal(size=(4, 3))

def my_matmul(a, b, my_sum, my_multiply):  # Proposed vectorized implementation
    tmp = my_multiply(a[..., None], b[None, ...])  # n×k(×1) ∘ (1×)k×m → n×k×m
    return my_sum(tmp)  # n×k×m → n×m

def my_matmul_loop(a, b, my_sum, my_multiply): # Explicit loopy implementation
    n_rows, n_cols = a.shape[0], b.shape[1]
    c = np.empty((n_rows, n_cols), dtype=float)
    for row in range(n_rows):
        for col in range(n_cols):
            tmp = my_multiply(a[row, :], b[:, col])
            # k → 1×k, as `my_sum` reduces axis 1; `squeeze()` for scalar result
            c[row, col] = my_sum(tmp[None, ...]).squeeze()
    return c

# Do standard sum and standard multiplication result in standard matrix product?
my_sum = lambda array: np.sum(array, axis=1)
my_multiply = np.multiply
assert np.allclose(a @ b, my_matmul_loop(a, b, my_sum, my_multiply))
assert np.allclose(a @ b, my_matmul(a, b, my_sum, my_multiply))

# Same result for loopy and vectorized version with custom sum and product?
my_sum = lambda array: np.max(array, axis=1)
my_multiply = lambda a, b: (a + 1) / (b - 1)
assert np.allclose(my_matmul_loop(a, b, my_sum, my_multiply),
                   my_matmul(a, b, my_sum, my_multiply))

Finally, to make sure that the given functions my_sum and my_multiply indeed behave as expected, we could add some sanity checks on the intermediate shapes:

def my_matmul(a, b, my_sum, my_multiply):  # Proposed vectorized implementation
    assert a.shape[1] == b.shape[0]
    tmp = my_multiply(a[..., None], b[None, ...])  # n×k(×1) ∘ (1×)k×m → n×k×m
    assert tmp.shape == (a.shape[0], a.shape[1], b.shape[1])
    tmp = my_sum(tmp)  # n×k×m → n×m
    assert tmp.shape == (a.shape[0], b.shape[1])
    return tmp