How to avoid padding summand in Tensorflow

user_2731113 0
Viewed 101

I have two tensors A, with shape (25, 1010, 7, 512), and B, with shape (10, 7).

I would like to add B to the first 10 slices using the second dimension of A in a way which could be done in Numpy as:

A[:, :10, :, :] += B[np.newaxis, :, :, np.newaxis]

(It's not important that the original tensor be updated, I'd be fine with either an in place update or a new tensor being created).

Of course this can't be done directly in tensorflow, and tf.tensor_scatter_nd_add doesn't seem to work for this case because the slices aren't into the final dimension of the tensor to be updated. Is there any way to accomplish this sparsely, or do I need to pad B out by a factor of 100?

EDIT: I messed up the numpy example when converting it from an in place update to an addition.

tensorflow2.0 tensorflow
1 Answers

EDIT: Since OP is trying to do an item assignment that is not available in TensorFlow (but available in PyTorch!!), you could use tf.concat to append over axis 1 the remaining tensor after the sum.

Please see my updated code -

import tensorflow as tf
#tf.__version__ = '2.3.1'

#numpy arrays with your shapes
a = np.random.random((25, 1010, 7, 512))
b = np.random.random((10, 7))

#convert to tensors
A = tf.convert_to_tensor(a)
B = tf.convert_to_tensor(b)

#Broadcasted sum >> concat with remaining array over axis 1
A = tf.concat([A[:,:10,:,:] + B[None, :, :, None], A[:,10:,:,:]], axis=1)
A.shape

TensorShape([25, 1010, 7, 512])

Just to check that you are getting the same answer as you would with item assignment in NumPy, you can do this -

#Item assigment sum over numpy arrays
a[:,:10,:,:] += b[None,:,:,None]

#Convert a to tensor and check against A from above.
tf.reduce_all(tf.convert_to_tensor(a) == A)

<tf.Tensor: shape=(), dtype=bool, numpy=True>
user_47559540
Related