使用tf.function重塑未知形状张量遇报错，求交换最后两维度方案

你的问题核心是Graph模式中不能直接迭代符号张量，原来的代码里[*tf.shape(tensor)[:-2]]试图把符号张量的切片结果拆成Python列表，这在@tf.function包装的函数里是不允许的。下面给两种优雅易读的解决方案：

方案一：用tf.transpose直接交换维度（推荐） #

交换维度用tf.transpose是最直观的，不管张量是4维还是5维，都能自动适配：

import tensorflow as tf
import logging

tensor = tf.random.uniform(shape=[4, 3, 2, 1])

@tf.function
def my_func():
    # 构造维度排列：前n-2个维度保持顺序，最后两个交换
    rank = tf.rank(tensor)
    perm = tf.concat([
        tf.range(rank - 2),  # 取前rank-2个维度的索引
        [rank - 1, rank - 2]  # 交换最后两个维度的索引
    ], axis=0)
    return tf.transpose(tensor, perm=perm)

logging.info(my_func())

方案二：用tf.reshape构造新形状 #

如果坚持要用reshape，需要用TensorFlow的张量拼接操作替代Python的拆包：

import tensorflow as tf
import logging

tensor = tf.random.uniform(shape=[4, 3, 2, 1])

@tf.function
def my_func():
    tensor_shape = tf.shape(tensor)
    # 拼接新形状：前n-2个形状 + 最后一个形状 + 倒数第二个形状
    new_shape = tf.concat([
        tensor_shape[:-2],
        tensor_shape[-1:],
        tensor_shape[-2:-1]
    ], axis=0)
    return tf.reshape(tensor, new_shape)

logging.info(my_func())

为什么原来的代码报错？ #

在@tf.function的Graph模式下，tf.shape(tensor)返回的是符号张量，它在图构建阶段没有具体数值。而[*tf.shape(tensor)[:-2]]这种写法需要迭代这个符号张量来拆成Python列表，TensorFlow的AutoGraph不支持这种操作，所以抛出了OperatorNotAllowedInGraphError。

内容的提问来源于stack exchange，提问作者Felix Schön