def row_times_matrix(row, matrix):
    return [sum(row[i] * matrix[i][j] for i in range(len(row)))
            for j in range(len(matrix[0]))]

x = [1, 2]
w1 = [[1, -1, 2, 0], [0, 1, -1, 3]]
w2 = [[1, 0], [0, 1], [2, 1], [1, -1]]
hidden = [max(0, v) for v in row_times_matrix(x, w1)]
dense = row_times_matrix(hidden, w2)
partials = []
for columns in ((0, 1), (2, 3)):
    local_w1 = [[row[j] for j in columns] for row in w1]
    local_hidden = [max(0, v) for v in row_times_matrix(x, local_w1)]
    local_w2 = [w2[j] for j in columns]
    partials.append(row_times_matrix(local_hidden, local_w2))
split = [sum(part[j] for part in partials) for j in range(2)]
assert split == dense == [7, -5]
print(f"dense={dense}; split={split}")
tokens, width, expanded, element_bytes = 128, 1024, 4096, 2
hidden_bytes = tokens * expanded * element_bytes
output_bytes = tokens * width * element_bytes
print(f"expanded_tensor_bytes={hidden_bytes}")
print(f"output_reduction_tensor_bytes={output_bytes}")
assert hidden_bytes == 4 * output_bytes