Assign part of a tensor to the corresponding places in another tensor

I could not find how to replace part of a tensor data with another tensor data. After digging a little bit, I see a lot reports that the tensor is not assignable data; some workarounds are suggested such as (https://github.com/tensorflow/tensorflow/issues/14132#issuecomment-483002522).

Let me give you a simple example of what I look for. I have two batches as follow:

x=·tf.random.uniform((2,3,2))
y= tf.random.uniform((2,3,2))
print (x)
print ('===================')
print (y)

The output of the two batches above as follow:

<tf.Tensor: shape=(2, 3, 2), dtype=float32, numpy=
array([[[0.17130184, 0.5413419 ],
        [0.6435019 , 0.82179713],
        [0.79388785, 0.9519701 ]],

       [[0.46769345, 0.9812336 ],
        [0.5796915 , 0.29866755],
        [0.0442245 , 0.86057484]]], dtype=float32)>
===================
<tf.Tensor: shape=(2, 3, 2), dtype=float32, numpy=
array([[[0.82299507, 0.8277409 ],
        [0.24233484, 0.4353037 ],
        [0.23145556, 0.00768614]],

       [[0.83972216, 0.03451204],
        [0.46768224, 0.44939125],
        [0.7840742 , 0.99360645]]], dtype=float32)>

I want to replace first row of each array in the x batch with the corresponding ones in the y batch.

I was expecting a result like :

<tf.Tensor: shape=(2, 3, 2), dtype=float32, numpy=
array([[[0.82299507, 0.8277409 ],  # copied from the y batch
        [0.6435019 , 0.82179713],
        [0.79388785, 0.9519701 ]], 

       [[0.83972216, 0.03451204],  # copied from the y batch
        [0.5796915 , 0.29866755],
        [0.0442245 , 0.86057484]]], dtype=float32)>

The following works when converting the batches to NumPy (but it is not what I want, I want to work with the tensors directly)

x = x.numpy()
y = y.numpy()

x[:, 0:1 , : ] = y[:, 0:1 , :]
x

The output is NumPy array which I can convert it to tensor again but I want to do such operation directly on the tensor.

array([[[0.82299507, 0.8277409 ],
        [0.6435019 , 0.82179713],
        [0.79388785, 0.9519701 ]],

       [[0.83972216, 0.03451204],
        [0.5796915 , 0.29866755],
        [0.0442245 , 0.86057484]]], dtype=float32)

Any help is much appreciated.

Solution 1:

Maybe just use tf.concat:

import tensorflow as tf
tf.random.set_seed(1)

x= tf.random.uniform((2,3,2))
y= tf.random.uniform((2,3,2))
print('X -->', x)
print('Y -->', y)
print('Z -->', tf.concat([y[:, 0:1 , :], x[:, 1: , : ]], axis=1))

X --> tf.Tensor(
[[[0.16513085 0.9014813 ]
  [0.6309742  0.4345461 ]
  [0.29193902 0.64250207]]

 [[0.9757855  0.43509948]
  [0.6601019  0.60489583]
  [0.6366315  0.6144488 ]]], shape=(2, 3, 2), dtype=float32)
Y --> tf.Tensor(
[[[0.51010704 0.44353175]
  [0.4085331  0.9924923 ]
  [0.68866396 0.34584963]]

 [[0.436067   0.601061  ]
  [0.45662427 0.75269794]
  [0.18799722 0.54875696]]], shape=(2, 3, 2), dtype=float32)
Z --> tf.Tensor(
[[[0.51010704 0.44353175]
  [0.6309742  0.4345461 ]
  [0.29193902 0.64250207]]

 [[0.436067   0.601061  ]
  [0.6601019  0.60489583]
  [0.6366315  0.6144488 ]]], shape=(2, 3, 2), dtype=float32)