Обновите значения в тензор по более глубоким измерениям

Question

У меня есть тензор A формы MxNxC, где M обозначает число примеров, N обозначает количество элементов, а C обозначает 3 угла поворота Эйлера.Также у меня есть тензор B аналогичной формы, но вместо углов есть координаты.Необходимо преобразовать оба этих тензора в один, содержащий аффинные матрицы преобразования, чтобы его форма была похожа на MxNx4x4.Я не знаю, как перебирать эти тензоры вместе, я искал tf.map_fn и tf.scan , но они повторяются только с первым измерением.Что мне нужно, так это какой-то метод применения функции, подобной приведенной ниже, ко всем элементам вдоль последних осей.

def f(angles, vector): #dimensions 3 or 3x1
    ...
    return matrix # dimension 4x4

Любая помощь будет полезна, спасибо!

Answer 1

вы можете попробовать что-то вроде этого:

A_flattened = tf.reshape(A, [-1, 3])# flatten it out
B_flattened = tf.reshape(B, [-1, 3])
AB_flattened = tf.map_fn(convert_to_mat, (A_flattened, B_flattened))# convert_to_mat should return a 4x4 matrix
AB = tf.reshape(AB_flattened, [M, N, 4, 4])

это должно сработать!