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-итерабельных объектов

С немедленным выполнением вы можете распаковать 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

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

Немедленные тензоры имеют конкретные значения, поэтому их форма всегда известна. Внутри 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. Но это должно быть постоянное значение.

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