Does (My)Conv is correct wrt NCHW

[jecampagne]

Hello,
I would like to be sure that I am doing correctly for my simple ConvNet dealing with MNIST (for exercice)
So, If get the images in a a batch stream that works like that

batches = data_stream()
x, y = next(batches) 
x.shape

(100, 1, 28, 28)

Should I have to do like that

from jax.example_libraries.stax import (Flatten, GeneralConv, Dense, MaxPool, BatchNorm, Relu, LogSoftmax)

MyConv = partial(GeneralConv, ('NCHW', 'HWIO', 'NCHW'))  # default Conv has 'NHWC',  'HWIO', 'NHWC'

init_fun, my_net = stax.serial(MyConv(32, (5, 5), (2, 2), padding="SAME"),
                                 BatchNorm(), Relu,
                                 MyConv(32, (5, 5), (2, 2), padding="SAME"),
                                 BatchNorm(), Relu,
                                 MyConv(10, (3, 3), (2, 2), padding="SAME"),
                                 BatchNorm(), Relu,
                                 MyConv(10, (3, 3), (2, 2), padding="SAME"), Relu,
                                 Flatten,                                 
                               Dense(num_labels),
                                 LogSoftmax)

Thanks.