Tensorflow 2D Matrices mutiplication returning list of matrix products

Question

I'm struggling with this problem in keras/tensorflow.
I'm implementing a user defined loss function and I have this problem: I have to multiply 2 matrices, obtaining a list of matrix products in the form
[column_0_matrix_1 x row_0_matrix_2], [column_1_matrix_1 x row_1_matrix_2] ecc.

Let's say I have

A = [[1 1]
     [3 2]]
B = [[4 1]
     [1 3]]

Then I want to have a list of products in the form

C = |[1] x [4 1]|, |[1] x [1 3]|
    |[3]        |  |[2]        |

Any idea? I tried by my self but always get back the product of the 2 starting matrices. Any help would by appreciated. Thank you

Answer 1

You could split each tensor and then use tf.linalg.matmul in a list comprehension to achieve what you want

import tensorflow as tf


a = tf.constant([[1, 1], [3, 2]])
b = tf.constant([[4, 1], [1, 3]])

a_split = tf.split(a, 2, 1)
b_split = tf.split(b, 2, 0)

[tf.linalg.matmul(x, y) for x, y in zip(a_split, b_split)]
# [<tf.Tensor: shape=(2, 2), dtype=int32, numpy=
#  array([[ 4,  1],
#         [12,  3]], dtype=int32)>,
# <tf.Tensor: shape=(2, 2), dtype=int32, numpy=
#  array([[1, 3],
#         [2, 6]], dtype=int32)>]