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.0, 2.0, 3.0]))
Traceback (most recent call last):

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

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

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