How to pad multiple tensors with one on main diagonal and zeros elsewhere?

Question

I have R as 2D rotation matrices of shape (N,2,2). Now I wish to extend each matrix to (3,3) 3D rotation matrices, i.e. to put zeros in each [:,:2,:2] and put 1 to [:,2,2].

How to do this in tensorflow?

UPDATE

I tried this way

    R = tf.get_variable(name='R', shape=np.shape(R_value), dtype=tf.float64,
                        initializer=tf.constant_initializer(R_value))
    eye = tf.eye(np.shape(R_value)[1]+1)
    right_column = eye[:2,2]
    bottom_row = eye[2,:]
    R = tf.concat([R, right_column], 3)
    R = tf.concat([R, bottom_row], 2)

but failed, because concat doesn't do broadcasting...

UPDATE 2

I made explicit broadcasting and also fixed wrong indices in concat calls:

    R = tf.get_variable(name='R', shape=np.shape(R_value), dtype=tf.float64,
                        initializer=tf.constant_initializer(R_value))
    eye = tf.eye(np.shape(R_value)[1]+1, dtype=tf.float64)

    right_column = eye[:2,2]
    right_column = tf.expand_dims(right_column, 0)
    right_column = tf.expand_dims(right_column, 2)
    right_column = tf.tile(right_column, (np.shape(R_value)[0], 1, 1))

    bottom_row = eye[2,:]
    bottom_row = tf.expand_dims(bottom_row, 0)
    bottom_row = tf.expand_dims(bottom_row, 0)
    bottom_row = tf.tile(bottom_row, (np.shape(R_value)[0], 1, 1))

    R = tf.concat([R, right_column], 2)
    R = tf.concat([R, bottom_row], 1)

The solutions looks rather complex. Are there any simpler ones?

Answer 1

first pad zeros to [N, 2, 2] to be [N, 3, 3] with padded = tf.pad(R, [[0, 0], [0, 1], [0, 1]])

then convert padded[N, 2, 2] to 1:

since tf.Tensor does not support assignment, you can do this with initializing a np.array, and then add them together.

arr = np.zeros((3, 3)) 
arr[2, 2] = 1
R = padded + arr # broadcast used here

now variable R is what you need