I’m trying to build a simple three-layer neural network with tensorflow. When I pass a singular input into the model, I get a valid output in the form I want. However, when I try to fit the data in order to train it, I get the following error:

2024-06-28 10:08:34.036743: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Incompatible shapes: [32,3] vs. [200,3]
         [[{{function_node __inference_one_step_on_data_575}}{{node sequential_1/rbf_layer_1/Sub}}]]

I have the code I’m using below.

Main Code:

import tensorflow as tf
import numpy as np
from RBF_layer import RBF_layer

# Generate source points inside the domain
def generateS(nS, R):
    dlong = np.pi * (3 - np.sqrt(5))
    long = 0
    dz = R * 2.0 / nS
    z = R - dz/2
    sourcenodes = np.zeros((3, int(nS)))
    for k in range(int(nS)):
        r = np.sqrt(R**2 - z**2)
        sourcenodes[1, k] = np.cos(long) * r
        sourcenodes[2, k] = np.sin(long) * r
        sourcenodes[0, k] = z
        z = z - dz
        long = long + dlong

    return sourcenodes

# Generate boundary points
def generateB(nB, part):
    dlong = np.pi * (3 - np.sqrt(5))
    long = 0
    dz = np.sqrt(1) * 2.0 / nB
    z = np.sqrt(1) - dz/2
    bdpt = np.zeros((3, int(part * nB)))
    for k in range(int(part * nB)):
        r = np.sqrt(1 - z**2)
        bdpt[1, k] = np.cos(long) * r
        bdpt[2, k] = np.sin(long) * r
        bdpt[0, k] = z
        z = z - dz
        long = long + dlong
    normalvec = bdpt

    return [bdpt, normalvec]

# Loss function for the neural network
def MSELoss(y_exact, y_pred):
    y_total = tf.subtract(y_pred, y_exact)
    return 0.5 * tf.matmul(y_total, tf.transpose(y_total))

Dirichlet = lambda x, y, z: x**2 + y**2 - 2* (z**2)

# Generate points used for the network as type float64
[bdpt, normalvec1] = generateB(200, 1)
bdpt1 = np.transpose(bdpt)
rad = 0.5
bdt = np.array(generateS(100, rad), dtype='float64')
testpt = np.transpose(bdt)
ds = 1.2
spt1 = np.transpose(generateS(100, ds))
rad2 = 0.2
spt2 = np.transpose(generateS(100, rad2))
spt = np.vstack((spt1, spt2), dtype='float64')
y_true =np.array(Dirichlet(testpt[:, 0], testpt[:, 1], testpt[:, 2]), dtype='float64')

#Build the network
model = tf.keras.models.Sequential()
model.add(RBF_layer(200, spt))
model.add(tf.keras.layers.Dense(units=1, activation="linear", kernel_initializer='LecunNormal', use_bias = False))
model.summary()

#Check the output for one test point
prediction = model(testpt[:1]).numpy()
tf.print(prediction)

#Build and train the network
model.compile(optimizer='adam', loss=MSELoss, metrics=['accuracy'])
model.fit(testpt, y_true, epochs=5)

Custom Layer Code:

import tensorflow as tf
import numpy as np
import math

class RBF_layer(tf.keras.layers.Layer):
    def __init__(self, num_outputs, source):
        #Build the layer
        super(RBF_layer, self).__init__()
        self.num_outputs = num_outputs
        self.source = source
    
    def build(self):
        self.bias = tf.Variable(initial_value=self.source, trainable=False, dtype='float64')
        
    def call(self, inputs):
        #Calculate the output given the input
        inputs = tf.cast(inputs, dtype=tf.dtypes.float64)
        r_0 = tf.math.subtract(inputs, self.bias)
        r = tf.sqrt(tf.reduce_sum(tf.pow(r_0,2), axis=1, keepdims=True))
        return tf.transpose(1/(4 * math.pi * r))
        

I’m a beginner in tensorflow and neural networks in general, so any help is appreciated as I don’t know how to begin tackling this error.