numpy - Naïve Tiled Matrix Multiplication with Python

Question

I'm trying to write a function to perform the matrix multiplication operation on two large matrices that are split up into smaller matrices (tiles), and I'm having difficulty conceptualizing a valid approach.

In the script below, two tensors/matrices, a and b, are segmented into multiple sub-tensors/matrices (tiles) of shape tile_shape. Given these tiles, I would like to perform operations that give the same result as a.dot(b), i.e., vector-matrix multiplication.

For simplicity sake, tiles are all the same shape, and are zero padded for cases where the length/width of a given tensor is not divisible by tile shape. I would like the approach to be modular, so that any given valid input shapes and tile shapes can be specified. Any help is greatly appreciated!

import numpy as np
import math

class Tiles:
    def __init__(self, tiles, tensor_shape):
        self.tiles = tiles
        self.tile_shape = tiles.shape[-2:]
        self.tensor_shape = tensor_shape

def gen_tiles(tensor, tile_shape):
    assert len(tensor.shape) == 2
    tiles = np.empty((math.ceil(tensor.shape[0] / tile_shape[0]), math.ceil(tensor.shape[1] / tile_shape[1]), tile_shape[0], tile_shape[1]))
    for i in range(0, tiles.shape[0]):
        for j in range(0, tiles.shape[1]):
            tile = tensor[i * tile_shape[0]:max((i + 1) * tile_shape[0], tile_shape[0] - 1), j * tile_shape[1]:max((j + 1) * tile_shape[1], tile_shape[1] - 1)]
            padded_tile = np.zeros(shape=tile_shape)
            padded_tile[:tile.shape[0],:tile.shape[1]] = tile
            tiles[i][j][:][:] = padded_tile

    return Tiles(tiles, tensor.shape)

def matmul_tiles(a, b):
    assert a.tensor_shape[1] == b.tensor_shape[0] and a.tile_shape == b.tile_shape
    result = np.zeros((a.tensor_shape[0], b.tensor_shape[1]))
    print(a.tiles.shape[0:2])
    print(b.tiles.shape[0:2])
    # Compute matrix multiplication using tiles

tile_shape = (2, 2)
a = np.random.randint(10, size=(2, 2))
b = np.random.randint(10, size=(2, 3))
a_tiles = gen_tiles(a, tile_shape)
b_tiles = gen_tiles(b, tile_shape)
product = a.dot(b)
tile_product = matmul_tiles(a_tiles, b_tiles)
assert np.isclose(product, tile_product)

question from:https://stackoverflow.com/questions/65622898/na%c3%afve-tiled-matrix-multiplication-with-python

Answer 1

Waitting for answers