Different outputs for windowed input on NN?

[StuartFarmer]

I have the following CNN written in Flax.

class Network(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=10, kernel_size=(4,), padding='SAME')(x)
        x = nn.tanh(x)
        x = x.reshape((x.shape[0], -1)) # flatten
        x = nn.Dense(1)(x)
        x = nn.tanh(x)
        
        return x

I initiate and pass a batch of my data through it and get my logits out. But, the logits I get out when I push a batch is different than when I push a single sample. For example:

x = Network()
variables = x.init(random.key(30000), jnp.ones(windows_train.shape))

x.apply(variables, np.expand_dims(windows_train[0], axis=0))
## Array([[-3.93813e-05]], dtype=float32)

x.apply(variables, np.expand_dims(windows_train[0:100][0], axis=0))
## Array([[-3.93813e-05]], dtype=float32)

x.apply(variables, windows_train[0:100])[0]
## Array([7.975578e-05], dtype=float32)

This is before training. The shape of the windows_train variable is (1337997, 100).

What am I doing wrong? Am I misunderstanding the library?