tf.unstack

Распаковывает заданное измерение тензора ранга R в тензоры ранга (R-1).

tf.unstack(
    value, num=None, axis=0, name='unstack'
)

Распаковывает тензоры из value, разбивая его по измерению axis.

x = tf.reshape(tf.range(12), (3,4))

p, q, r = tf.unstack(x)
p.shape.as_list()
[4]

i, j, k, l = tf.unstack(x, axis=1)
i.shape.as_list()
[3]

Это противоположно функции stack.

x = tf.stack([i, j, k, l], axis=1)

Более общо, если у вас есть тензор формы (A, B, C, D).

A, B, C, D = [2, 3, 4, 5]
t = tf.random.normal(shape=[A, B, C, D])

Количество возвращаемых тензоров равно длине целевого axis:

axis = 2
items = tf.unstack(t, axis=axis)
len(items) == t.shape[axis]
True

Форма каждого результирующего тензора равна форме входного тензора, при этом целевое измерение axis удалено.

items[0].shape.as_list()  # [A, B, D]
[2, 3, 5]

Значение каждого тензора items[i] равно срезу input вдоль axis в индексе i.

for i in range(len(items)):
  slice = t[:,:,i,:]
  assert tf.reduce_all(slice == items[i])

Распаковка Python-итерируемых объектов

При использовании режима выполнения Eager вы можете распаковать 0-й ось тензора с помощью Python-итерируемой распаковки:

t = tf.constant([1,2,3])
a,b,c = t

unstack все еще необходимо, потому что распаковка итерируемых объектов не работает в @tf.function: Символьные тензоры не итерируемы.

Здесь необходимо использовать tf.unstack:

@tf.function
def bad(t):
  a,b,c = t
  return a

bad(t)
Traceback (most recent call last):

OperatorNotAllowedInGraphError: ...

@tf.function
def good(t):
  a,b,c = tf.unstack(t)
  return a

good(t).numpy()
1

Неизвестные формы

Тензоры Eager имеют конкретные значения, поэтому их форма всегда известна. Внутри tf.function символьные тензоры могут иметь неизвестные формы. Если длина axis неизвестна, tf.unstack завершится ошибкой, так как она не может обрабатывать неизвестное количество тензоров:

@tf.function(input_signature=[tf.TensorSpec([None], tf.float32)])
def bad(t):
  tensors = tf.unstack(t)
  return tensors[0]

bad(tf.constant([1,2,3]))
Traceback (most recent call last):

ValueError: Cannot infer argument `num` from shape (None,)

Если вам известна длина axis, вы можете передать ее в качестве аргумента num. Но это должно быть константное значение.

Если вам действительно требуется переменное количество тензоров в одном трассировочном следе tf.function, вам придется использовать явные циклы и tf.TensorArray вместо этого.