python - Combine arbitrary shaped tensors

Question

I'd like to combine two variable length tensors.

Since they don't match in shape I can't use tf.concat or tf.stack.

So I thought I'd flatten one and then append it to each element of the other - but I don't see how to do that.

For example,

a = [ [1,2], [3,4] ]
flat_b = [5, 6]

combine(a, flat_b) would be [ [ [1,5,6], [2,5,6] ],
                              [ [3,5,6], [4,5,6] ] ]

Is there a method like this?

question from:https://stackoverflow.com/questions/65649660/combine-arbitrary-shaped-tensors

Answer 1

Using tf.map_fn with tf.concat, Example code:

import tensorflow as tf

a = tf.constant([ [1,2], [3,4] ])
flat_b = [5, 6]
flat_a = tf.reshape(a, (tf.reduce_prod(a.shape).numpy(), ))[:, tf.newaxis]
print(flat_a)
c = tf.map_fn(fn=lambda t: tf.concat([t, flat_b], axis=0), elems=flat_a)
c = tf.reshape(c, (-1, a.shape[1], c.shape[1]))
print(c)

Outputs:

tf.Tensor(
[[1]
 [2]
 [3]
 [4]], shape=(4, 1), dtype=int32)
tf.Tensor(
[[[1 5 6]
  [2 5 6]]

 [[3 5 6]
  [4 5 6]]], shape=(2, 2, 3), dtype=int32)