Tensorflowjs expects 3 Dimensions, even though I trained it with 2 in Python

Question

I am trying to test a model which I trained using XOR-Operator values in Python:

x_train = [[1,0],[0,1],[0,0],[1,1]]
x_label = [1,1,0,0]

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(2,1)),
    keras.layers.Dense(8, activation="relu"),
    keras.layers.Dense(2, activation="softmax")
])

model.compile(optimizer="adam", loss="sparse_categorical_crossentropy",metrics=["accuracy"])
model.fit(x_train, x_label, epochs=15)

prediction = model.predict([[5,5]])
print("Prediction: ", np.argmax(prediction[0]))

It works great in Python, but when I transform it into a Tensorflow JS file (model.json), it doesn't work anymore:

tfjs.converters.save_keras_model(model, "test_model")

As you can see above, I trained it with input_shape=(2,1). But when I try it in Nodejs, I get the error:

ValueError: Error when checking : expected flatten_input to have 3 dimension(s), but got array with shape [2,1]

My JS Code:

tf.loadLayersModel(model_url).then((model)=>{
        const input = tf.tensor([1,1],[2,1], "int32");
        const prediciton = model.predict(input)
        console.log(prediciton);
        response.send(prediciton);
        return null;
    }).catch((e)=>{
        console.log(e);
    });

Thank you!

Answer 1

tf.tensor([1,1],[2,1]) (equivalent to the following array: [[1], [1]] will not work even in python). In python when the shape does not match, the dimension is expanded on the last axis. The resulting tensor will thus have the shape 2, 1, 1 which does not match the input shape that is expecting an input of shape [b, 2, 1].

However tf.tensor([1,1],[1,2]) will work in python because of the above. It seems however than in js broadcasting is only done when the input tensor is a tensor1d. So even this will not work.

To solve the issue in js, you can just consider expanding the tensor on the first axis (axis = 0) with the following:

tf.tensor([1,1],[1, 2,1], "int32")