tf.split

Разделяет тензор value на список подтензоров.

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

См. также tf.unstack.

Если num_or_size_splits является целым числом, то value разбивается вдоль измерения axis на num_or_size_splits меньших тензоров. Это требует, чтобы value.shape[axis] было кратно num_or_size_splits.

Если num_or_size_splits является одномерным тензором (или списком), то value разбивается на len(num_or_size_splits) элементов. Форма i-го элемента имеет тот же размер, что и value, за исключением измерения axis, где размер равен num_or_size_splits[i].

Пример:

x = tf.Variable(tf.random.uniform([5, 30], -1, 1))

# Split `x` into 3 tensors along dimension 1
s0, s1, s2 = tf.split(x, num_or_size_splits=3, axis=1)
tf.shape(s0).numpy()
array([ 5, 10], dtype=int32)

# Split `x` into 3 tensors with sizes [4, 15, 11] along dimension 1
split0, split1, split2 = tf.split(x, [4, 15, 11], 1)
tf.shape(split0).numpy()
array([5, 4], dtype=int32)
tf.shape(split1).numpy()
array([ 5, 15], dtype=int32)
tf.shape(split2).numpy()
array([ 5, 11], dtype=int32)