如何使用tensorflow神经网络对对称函数进行回归建模

import tensorflow as tf
import numpy as np

model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(input_dim,)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(output_dim)
])

model.compile(optimizer='adam', loss='mse', metrics=['mae'])

model.fit(X_train, Y_train, epochs=100, batch_size=32, validation_data=(X_val, Y_val))

loss, mae = model.evaluate(X_test, Y_test)

predictions = model.predict(X_new)