tf.split

Разделяет тензор на подтензоры.

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

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

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

Например:

# 'value' is a tensor with shape [5, 30]
# Split 'value' into 3 tensors with sizes [4, 15, 11] along dimension 1
split0, split1, split2 = tf.split(value, [4, 15, 11], 1)
tf.shape(split0)  # [5, 4]
tf.shape(split1)  # [5, 15]
tf.shape(split2)  # [5, 11]
# Split 'value' into 3 tensors along dimension 1
split0, split1, split2 = tf.split(value, num_or_size_splits=3, axis=1)
tf.shape(split0)  # [5, 10]