如何在tensorflow中有效地对张量进行分组

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

import tensorflow as tf

# 创建一个形状为[6, 6]的张量
x = tf.constant([[1, 2, 3, 4, 5, 6],
                 [7, 8, 9, 10, 11, 12],
                 [13, 14, 15, 16, 17, 18],
                 [19, 20, 21, 22, 23, 24],
                 [25, 26, 27, 28, 29, 30],
                 [31, 32, 33, 34, 35, 36]])

# 在第一个维度上将张量分割成3个子张量
result = tf.split(x, num_or_size_splits=3, axis=0)

print(result)

[<tf.Tensor: shape=(2, 6), dtype=int32, numpy=
array([[1, 2, 3, 4, 5, 6],
       [7, 8, 9, 10, 11, 12]], dtype=int32)>,
 <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
array([[13, 14, 15, 16, 17, 18],
       [19, 20, 21, 22, 23, 24]], dtype=int32)>,
 <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
array([[25, 26, 27, 28, 29, 30],
       [31, 32, 33, 34, 35, 36]], dtype=int32)>]