I am interested in implementing this paper on Kronecker Recurrent Units in TensorFlow.
This involves the computation of a Kronecker Product. TensorFlow does not have an operation for Kronecker Products. I am looking for an efficient and robust way to compute this.
Does this exist, or would I need to define a TensorFlow op manually?
TensorFlow 1.7+ provides the function kronecker_product in tf.contrib.kfac.utils.kronecker_product:
a = tf.eye(3)
b = tf.constant([[1., 2.], [3., 4.]])
kron = tf.contrib.kfac.utils.kronecker_product(a, b)
tf.Session().run(kron)
Output:
array([[1., 2., 0., 0., 0., 0.],
[3., 4., 0., 0., 0., 0.],
[0., 0., 1., 2., 0., 0.],
[0., 0., 3., 4., 0., 0.],
[0., 0., 0., 0., 1., 2.],
[0., 0., 0., 0., 3., 4.]], dtype=float32)
How about something like this:
def kron(x, y):
"""Computes the Kronecker product of two matrices.
Args:
x: A matrix (or batch thereof) of size m x n.
y: A matrix (or batch thereof) of size p x q.
Returns:
z: Kronecker product of matrices x and y of size mp x nq
"""
with tf.name_scope('kron'):
x = tf.convert_to_tensor(x, dtype_hint=tf.float32)
y = tf.convert_to_tensor(y, dtype_hint=x.dtype)
def _maybe_expand(x):
xs = tf.pad(
tf.shape(x),
paddings=[[tf.maximum(2 - tf.rank(x), 0), 0]],
constant_values=1)
x = tf.reshape(x, xs)
_, mx, nx = tf.split(xs, num_or_size_splits=[-1, 1, 1])
return x, mx, nx
x, mx, nx = _maybe_expand(x)
y, my, ny = _maybe_expand(y)
x = x[..., :, tf.newaxis, :, tf.newaxis]
y = y[..., tf.newaxis, :, tf.newaxis, :]
z = x * y
bz = tf.shape(z)[:-4]
z = tf.reshape(z, tf.concat([bz, mx * my, nx * ny], axis=0))
return z
This solution: