What is the correct way to pass TensorFlow arrays to Sympy's lambdify() function?

Question

I am trying to generate a function from Sympy using lambdify() with tensorflow module. I am not getting the right answer when I pass a tensorflow array. I think this could be a problem with how I am passing the arguments to the function.

import sympy as sym
import tensorflow as tf

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

# Construct NN function
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"))

# Domain
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)

# Test Sympy Function
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])]

The output of res1 is wrong. The output of res2 is what I want, but do not want to use list comprehension.

The function F above should generate a sequence of scalar values, as it is in the variable res2.

res1 is sequence of a two element vector, with incorrect values.

Any help in resolving this will be greatly appreciated.

Answer 1

help(F_func) displays

Help on function _lambdifygenerated:

_lambdifygenerated(_Dummy_23, _Dummy_22)
    Created with lambdify. Signature:
    
    func(x, f)
    
    Expression:
    
    x1 + x2 + f(x1, x2)
    
    Source code:
    
    def _lambdifygenerated(_Dummy_23, _Dummy_22):
        [x1, x2] = _Dummy_23
        return _Dummy_22 + x1 + x2

The F sympy expression is:

In [8]: F
Out[8]: 
x1 + x2 + f(x1, x2)

With numpy inputs, the calculation is straight forward, add the X terms, and add on the f term:

In [12]: F_func([np.arange(3),3.2],1.23)
Out[12]: array([4.43, 5.43, 6.43])

There's no iteration implied in lambdify.