What is the most efficient way to compute a Kronecker Product in TensorFlow?

Question

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?

Answer 1

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)