如何正确将TensorFlow数组传入Sympy的lambdify()函数？

尝试使用Sympy的lambdify()生成TensorFlow兼容函数时，批量传入TensorFlow数组得到错误结果：res1输出包含两个元素的向量序列，而通过列表推导得到的res2是预期的标量序列，但希望找到无需循环的批量传参方式。

复现代码 #

import sympy as sym
import tensorflow as tf

from keras.models import Sequential
from keras.layers import Dense

# 构建神经网络函数
fhat = Sequential()
fhat.add(Dense(20, input_dim=2))
fhat.add(Dense(20, activation="tanh"))
fhat.add(Dense(20, activation="tanh"))
fhat.add(Dense(20, activation="tanh"))
fhat.add(Dense(20, activation="tanh"))
fhat.add(Dense(20, activation="tanh"))
fhat.add(Dense(1,activation="sigmoid"))

# 定义域
lim = tf.cast(5,dtype=tf.float32)
N = 5
nSamples = N**2

xx = tf.linspace(-lim,lim,N)
X1,X2 = tf.meshgrid(xx,xx)

XX1 = tf.reshape(X1,[nSamples])
XX2 = tf.reshape(X2,[nSamples])
X = tf.stack([XX1,XX2],axis=1)

# 测试Sympy函数
nx = 2;
x = sym.symbols(f'x1:{nx+1}', real=True)
f = sym.Function('f',real=True)(*x)

F = sum(x) + f
F_func = sym.lambdify([x,f], F, 'tensorflow')

Fhat = fhat(X)
res1 = F_func(X,Fhat)
res2 = [F_func(X[i],Fhat[i]) for i in range(XX1.shape[0])]

问题原因 #

x是Sympy的符号元组(x1, x2)，当将形状为(nSamples,2)的X传入F_func时，lambdify会把整个X当作元组x的输入，导致sum(x)对X的列求和（每个样本得到两个元素的和），再加上Fhat的结果后，最终输出每个样本对应两个元素的向量，与预期的标量序列不符。

解决方案 #

方法1：拆分Sympy符号参数

将x拆分为单独的符号变量，调整lambdify的参数列表，调用时传入X的每一列：

# 修改Sympy部分的定义
nx = 2
x1, x2 = sym.symbols('x1 x2', real=True)
f = sym.Function('f', real=True)(x1, x2)

F = x1 + x2 + f
# 参数列表改为单独的符号变量
F_func = sym.lambdify([x1, x2, f], F, 'tensorflow')

# 调用时拆分X的两列作为输入
res1 = F_func(X[:, 0], X[:, 1], Fhat)

方法2：使用TensorFlow map_fn批量处理

保持原有Sympy定义，用tf.map_fn实现向量化的批量处理，避免显式循环：

# 用tf.map_fn替代列表推导，自动批量处理每个样本
res1 = tf.map_fn(
    lambda args: F_func(args[0], args[1]),
    (X, Fhat),
    dtype=tf.float32
)

说明 #

• 方法1更贴合lambdify的参数传递逻辑，直接对应每个Sympy符号传入张量，性能最优。
• 方法2适用于无法修改Sympy定义的场景，利用TensorFlow原生的批量操作实现，避免Python层循环。

内容的提问来源于stack exchange，提问作者raktimb