I was trying to build a simple example of NN using TensorFlow.js, with source converted from an other old TensorFlow (Python) project of mine.
But for some reason, the model fails to interpret the input as I wanted, which is (2, 3, 3) for input, and (2) for output

target expected a batch of elements where each example has shape [2]  
(i.e.,tensor shape [*,2]) but the target received an input with 2 examples,  
each with shape [1] (tensor shape [2,1])

Full code:

  // define model
  let model = new Sequential({
    name: "test",
    layers: [
      layers.flatten({ inputShape: [3, 3] }),
      layers.dense({ units: 9, activation: 'relu' }),
      layers.dense({ units: 2 })
    ]
  });
  model.compile({
    optimizer: 'Adam',
    // should be loss=kr.losses.SparseCategoricalCrossentropy(from_logits=True)
    // as per the old code, but I don't know how yet
    // and actually, I'm not sure there would be any difference
    loss: losses.meanSquaredError, 
    metrics: ['accuracy']
  })

  // define input and output
  let X = tensor(
    [
      [
        [0, 1, 0],
        [1, 0, 1],
        [0, 1, 0]
      ],
      [
        [1, 0, 1],
        [0, 1, 0],
        [1, 0, 1]
      ],
    ]
  );
  let Y = tensor([1, 0]);

  model.fit(X, Y, { epochs: 1000 })
    .then(v => console.log("success", v))
    .catch((e: Error) => console.log(e.message)); // <- message logged from here

However, as I described, the shape of X in this case would be [2, 3, 3], not [2]:

console.log(X.shape);
> (3) [2, 3, 3]

I'm stuck at this one for several days now. Where did [*,2] even come from?