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?

Seduce

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)

Vestige

This function no longer seems to be available in Tensorflow 1.11. – Pulsatile

Answer 2

If you will read the math definition of conv2d_transpose and see what Kronecker product calculates, you will see that with the appropriate size of stides for conv2d_tranpose (width, height of the second matrix), it does the same thing.

Moreover you even have batching of Kronecker's product out of the box with conv2d_transpose.

---

Here is an example of you which calculates the Kronecker's product for matrices from wiki.

import tensorflow as tf
a = [[1, 2], [3, 4]]
b = [[0, 5], [6, 7]]

i, k, s = len(a), len(b), len(b)
o = s * (i - 1) + k

a_tf  = tf.reshape(tf.constant(a, dtype=tf.float32), [1, i, i, 1])
b_tf = tf.reshape(tf.constant(b, dtype=tf.float32), [k, k, 1, 1])

res = tf.squeeze(tf.nn.conv2d_transpose(a_tf, b_tf, (1, o, o, 1), [1, s, s, 1], "VALID"))

with tf.Session() as sess:
    print sess.run(res)

Notice that in the case of a non-square matrix, you will need to calulcate more dimensions in the lines:

i, k, s = len(a), len(b), len(b)
o = s * (i - 1) + k

and use them properly as your strides/outputs arguments.

Zeidman

why not directly 'o=i*s' – Liuka

Answer 3

Here's the utility I use for this. See kronecker_test for example of usage

def fix_shape(tf_shape):
  return tuple(int(dim) for dim in tf_shape)

def concat_blocks(blocks, validate_dims=True):
  """Takes 2d grid of blocks representing matrices and concatenates to single
  matrix (aka ArrayFlatten)"""

  if validate_dims:
    col_dims = np.array([[int(b.shape[1]) for b in row] for row in blocks])
    col_sums = col_dims.sum(1)
    assert (col_sums[0] == col_sums).all()
    row_dims = np.array([[int(b.shape[0]) for b in row] for row in blocks])
    row_sums = row_dims.sum(0)
    assert (row_sums[0] == row_sums).all()

  block_rows = [tf.concat(row, axis=1) for row in blocks]
  return tf.concat(block_rows, axis=0)

def chunks(l, n):
  """Yield successive n-sized chunks from l."""
  for i in range(0, len(l), n):
    yield l[i:i + n]

from tensorflow.python.framework import ops
original_shape_func = ops.set_shapes_for_outputs
def disable_shape_inference():
  ops.set_shapes_for_outputs = lambda _: _

def enable_shape_inference():
  ops.set_shapes_for_outputs = original_shape_func

def kronecker(A, B, do_shape_inference=True):
  """Kronecker product of A,B.
  turn_off_shape_inference: if True, makes 10x10 kron go 2.4 sec -> 0.9 sec
  """

  Arows, Acols = fix_shape(A.shape)
  Brows, Bcols = fix_shape(B.shape)
  Crows, Ccols = Arows*Brows, Acols*Bcols

  temp = tf.reshape(A, [-1, 1, 1])*tf.expand_dims(B, 0)
  Bshape = tf.constant((Brows, Bcols))

  # turn off shape inference
  if not do_shape_inference:
    disable_shape_inference()

  # [1, n, m] => [n, m]
  slices = [tf.reshape(s, Bshape) for s in tf.split(temp, Crows)]

  #  import pdb; pdb.set_trace()
  grid = list(chunks(slices, Acols))
  assert len(grid) == Arows
  result = concat_blocks(grid, validate_dims=do_shape_inference)

  if not do_shape_inference:
    enable_shape_inference()
    result.set_shape((Arows*Brows, Acols*Bcols))

  return result

def kronecker_test():
  A0 = [[1,2],[3,4]]
  B0 = [[6,7],[8,9]]
  A = tf.constant(A0)
  B = tf.constant(B0)
  C = kronecker(A, B)
  sess = tf.Session()
  C0 = sess.run(C)
  Ct = [[6, 7, 12, 14], [8, 9, 16, 18], [18, 21, 24, 28], [24, 27, 32, 36]]
  Cnp = np.kron(A0, B0)
  check_equal(C0, Ct)
  check_equal(C0, Cnp)

Guardhouse

Answer 4

Try the following solution, see if it works for you:

def tf_kron(a,b):
    a_shape = [a.shape[0].value,a.shape[1].value]
    b_shape = [b.shape[0].value,b.shape[1].value]
    return tf.reshape(tf.reshape(a,[a_shape[0],1,a_shape[1],1])*tf.reshape(b,[1,b_shape[0],1,b_shape[1]]),[a_shape[0]*b_shape[0],a_shape[1]*b_shape[1]])

Deshabille

Answer 5

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:

• supports batches
• supports broadcasting
• works in xla
• clearly shows the relationship between numpy broadcasting and kronecker products.

Binns