bug: Equivalent code works in tensorflow python but not tensorflow-js

TensorFlow.js version

Tried in versions 1.1.2 and 1.2.5

Describe the problem or feature request

The same network works in python version of tensorflow but the equivalent code does not work with tensorflow-js resulting in the error Error: Size(9) must match the product of shape 1,1,1,1

Code to reproduce the bug / link to feature request

tensorflow-js version

const tf = require('@tensorflow/tfjs');
require('@tensorflow/tfjs-node');  // Use '@tensorflow/tfjs-node-gpu' if running with GPU.
const { regularizers } = tf;
const { conv2d, batchNormalization, leakyReLU, input, dense, flatten } = tf.layers;

let main_input = input({ shape: [2, 3, 3], name: 'main_input' })

let x = conv2d({
  kernelSize: [4, 4],
  filters: 75,
  activation: 'linear',
  dataFormat: "channelsFirst",
  padding: 'same',
  useBias: false,
  kernelRegularizer: regularizers.l2(0.0001)
}).apply(main_input)

x = batchNormalization({ axis: 1 }).apply(x)
x = leakyReLU().apply(x)

x = conv2d({
  filters: 1,
  kernelSize: [1, 1],
  activation: 'linear',
  dataFormat: "channelsFirst",
  padding: 'same',
  name: 'value_head_conv2d',
  useBias: false,
  kernelRegularizer: regularizers.l2(0.0001)
}).apply(x)

x = batchNormalization({ axis: 1, name: 'value_head_batch_normalization' }).apply(x)
x = leakyReLU({ name: 'value_head_leaky_relu_1' }).apply(x)

x = flatten({ name: 'value_head_flatten' }).apply(x)

x = dense({
  units: 1,
  useBias: false,
  activation: 'tanh',
  kernelRegularizer: regularizers.l2(0.0001),
  name: 'output_1'
}).apply(x)

let model = tf.model({ inputs: [main_input], outputs: [x] })
model.compile({
  loss: { 'output_1': 'meanSquaredError' },
  optimizer: tf.train.sgd(0.1),
})

let model_input = tf.tensor([[[[1,0,0],[1,1,0],[0,0,0]],[[0,0,1],[0,0,0],[1,1,0]]]])
let model_output = tf.tensor([[1]])

model.fit(model_input, { 'output_1': model_output }, {
  epochs: 10,
  callbacks: {
    onEpochEnd: async (epoch, log) => {
      console.log(`Epoch ${epoch}: loss = ${log.loss}`);
    }
  }})

python version

import tensorflow as tf

from keras.models import Sequential, load_model, Model
from keras.layers import Input, Dense, Conv2D, Flatten, BatchNormalization, Activation, LeakyReLU, add
from keras.optimizers import SGD
from keras import regularizers

import numpy as np

main_input = Input(shape = [2, 3, 3], name = 'main_input')

x = Conv2D(
		filters = 75
		, kernel_size = (4,4)
		, data_format="channels_first"
		, padding = 'same'
		, use_bias=False
		, activation='linear'
		, kernel_regularizer = regularizers.l2(0.0001)
		)(main_input)

x = BatchNormalization(axis=1)(x)
x = LeakyReLU()(x)

x = Conv2D(
		filters = 1
		, kernel_size = (1,1)
		, data_format="channels_first"
		, padding = 'same'
		, use_bias=False
		, activation='linear'
		, kernel_regularizer = regularizers.l2(0.0001)
		)(x)


x = BatchNormalization(axis=1)(x)
x = LeakyReLU()(x)

x = Flatten()(x)

x = Dense(
  1
  , use_bias=False
  , activation='tanh'
  , kernel_regularizer=regularizers.l2(0.0001)
  , name = 'output_1'
  )(x)

model = Model(inputs=[main_input], outputs=[x])

model.compile(loss={'output_1': 'mean_squared_error'},
			optimizer=SGD(lr=0.001, momentum = 0.001)	
			)

model_input = np.array([[[[1,0,0],[1,1,0],[0,0,0]],[[0,0,1],[0,0,0],[1,1,0]]]])

model_output = np.array([[1]])

model.fit(model_input, {'output_1': model_output}, epochs= 10)

results

tensorflowjs result: Error: Size(9) must match the product of shape 1,1,1,1 python result: runs training as expected, no errors.

[caisq]

@lp74 Please construct the regularizer with an object argument: {l2: 0.01}, instad of 0.01.

[rthadur]

This question is better asked on StackOverflow since it is not a bug or feature request. There is also a larger community that reads questions there.