Modifying (keras/tensorflow) Tensors using numpy methods

Here is a way to do that with TensorFlow:

import tensorflow as tf

data = tf.placeholder(tf.float32, [None, None])
n = tf.placeholder(tf.int32, [])
eye = tf.eye(n)
mult = data[:, tf.newaxis, :, tf.newaxis] * eye[tf.newaxis, :, tf.newaxis, :]
result = tf.reshape(mult, n * tf.shape(data))
with tf.Session() as sess:
    a = sess.run(result, feed_dict={data: [[1, 2], [3, 4]], n: 3})
    print(a)

Output:

[[1. 0. 0. 2. 0. 0.]
 [0. 1. 0. 0. 2. 0.]
 [0. 0. 1. 0. 0. 2.]
 [3. 0. 0. 4. 0. 0.]
 [0. 3. 0. 0. 4. 0.]
 [0. 0. 3. 0. 0. 4.]]

By the way, you can do basically the same in NumPy, which should be faster than your current solution:

import numpy as np

data = np.array([[1, 2], [3, 4]])
n = 3
eye = np.eye(n)
mult = data[:, np.newaxis, :, np.newaxis] * eye[np.newaxis, :, np.newaxis, :]
result = np.reshape(mult, (n * data.shape[0], n * data.shape[1]))
print(result)
# The output is the same as above

EDIT:

I'll try to give some intuition about why/how this works, sorry if it's too long. It is not that hard but I think it's sort of tricky to explain. Maybe it is easier to see how the following multiplication works

import numpy as np

data = np.array([[1, 2], [3, 4]])
n = 3
eye = np.eye(n)
mult1 = data[:, :, np.newaxis, np.newaxis] * eye[np.newaxis, np.newaxis, :, :]

Now, mult1 is a sort of "matrix of matrices". If I give two indices, I will get the diagonal matrix for the corresponding element in the original one:

print(mult1[0, 0])
# [[1. 0. 0.]
#  [0. 1. 0.]
#  [0. 0. 1.]]

So you could say this matrix could be visualize like this:

| 1 0 0 |  | 2 0 0 |
| 0 1 0 |  | 0 2 0 |
| 0 0 1 |  | 0 0 2 |

| 3 0 0 |  | 4 0 0 |
| 0 3 0 |  | 0 4 0 |
| 0 0 3 |  | 0 0 4 |

However this is deceiving, because if you try to reshape this to the final shape the result is not the right one:

print(np.reshape(mult1, (n * data.shape[0], n * data.shape[1])))
# [[1. 0. 0. 0. 1. 0.]
#  [0. 0. 1. 2. 0. 0.]
#  [0. 2. 0. 0. 0. 2.]
#  [3. 0. 0. 0. 3. 0.]
#  [0. 0. 3. 4. 0. 0.]
#  [0. 4. 0. 0. 0. 4.]]

The reason is that reshaping (conceptually) "flattens" the array first and then gives the new shape. But the flattened array in this case is not what you need:

print(mult1.ravel())
# [1. 0. 0. 0. 1. 0. 0. 0. 1. 2. 0. 0. 0. 2. 0. ...

You see, it first traverses the first submatrix, then the second, etc. What you want though is for it to traverse first the first row of the first submatrix, then the first row of the second submatrix, then second row of first submatrix, etc. So basically you want something like:

• Take the first two submatrices (the ones with 1 and 2)
  • Take all the first rows ([1, 0, 0] and [2, 0, 0]).
    • Take the first of these ([1, 0, 0])
      • Take each of its elements (1, 0 and 0).

And then continue for the rest. So if you think about it, we traversing first the axis 0 (row of "matrix of matrices"), then 2 (rows of each submatrix), then 1 (column of "matrix of matrices") and finally 3 (columns of submatrices). So we can just reorder the axis to do that:

mult2 = mult1.transpose((0, 2, 1, 3))
print(np.reshape(mult2, (n * data.shape[0], n * data.shape[1])))
# [[1. 0. 0. 2. 0. 0.]
#  [0. 1. 0. 0. 2. 0.]
#  [0. 0. 1. 0. 0. 2.]
#  [3. 0. 0. 4. 0. 0.]
#  [0. 3. 0. 0. 4. 0.]
#  [0. 0. 3. 0. 0. 4.]]

And it works! So in the solution I posted, to avoid the tranposing, I just make the multiplication so the order of the axes is exactly that:

mult = data[
        :,           # Matrix-of-matrices rows
        np.newaxis,  # Submatrix rows
        :,           # Matrix-of-matrices columns
        np.newaxis   # Submatrix columns
    ] * eye[
        np.newaxis,  # Matrix-of-matrices rows
        :,           # Submatrix rows
        np.newaxis,  # Matrix-of-matrices columns
        :            # Submatrix columns
    ]

I hope that makes it slightly clearer. To be honest, in this case in particular I could came up with the solution quickly because I had to solve a similar problem not too long ago, and I guess you end up building an intuition of these things.

Another way to achieve the same effect in numpy is to use the following:

A = np.array([[1,2],
            [3,4]])
B = np.repeat(np.repeat(A, 3, axis=0), 3, axis=1) * np.tile(np.eye(3), (2,2))

Then, to replicate it in tensorflow, we can use tf.tile, but there is no tf.repeat, however someone has provided this function on tensorflow tracker.

def tf_repeat(tensor, repeats):
    """
    Args:

    input: A Tensor. 1-D or higher.
    repeats: A list. Number of repeat for each dimension, length must be the same as the number of dimensions in input

    Returns:

    A Tensor. Has the same type as input. Has the shape of tensor.shape * repeats
    """
    with tf.variable_scope("repeat"):
        expanded_tensor = tf.expand_dims(tensor, -1)
        multiples = [1] + list(repeats)
        tiled_tensor = tf.tile(expanded_tensor, multiples=multiples)
        repeated_tesnor = tf.reshape(tiled_tensor, tf.shape(tensor) * repeats)
    return repeated_tesnor

and thus the tensorflow implementation will look like the following. Here I also consider that the first dimension represents batches, and thus we do not operate on it.

N = 2
M = 3
nbatch = 2
Ain = np.reshape(np.arange(1, 1 + N*N*nbatch), (nbatch, N, N))

A = tf.placeholder(tf.float32, shape=(nbatch, N, N))
B = tf.tile(tf.eye(M), [N, N]) * tf_repeat(A, [1, M, M])

with tf.Session() as sess:
    print(sess.run(C, feed_dict={A: Ain}))

and the result:

 [[[1. 0. 0. 2. 0. 0.]
  [0. 1. 0. 0. 2. 0.]
  [0. 0. 1. 0. 0. 2.]
  [3. 0. 0. 4. 0. 0.]
  [0. 3. 0. 0. 4. 0.]
  [0. 0. 3. 0. 0. 4.]]

 [[5. 0. 0. 6. 0. 0.]
  [0. 5. 0. 0. 6. 0.]
  [0. 0. 5. 0. 0. 6.]
  [7. 0. 0. 8. 0. 0.]
  [0. 7. 0. 0. 8. 0.]
  [0. 0. 7. 0. 0. 8.]]]