mnist example flatten

[deadsoul44]

flax/examples/mnist/train.py

Line 47 in 13fbace

x = x.reshape((x.shape[0], -1)) # flatten

The line above claims to flatten the array but it does nothing.

x = jnp.ones((2, 3))
print(x)
# prints:
[[1. 1. 1.]
 [1. 1. 1.]]

y = x.reshape((x.shape[0], -1))
print(y)
# prints:
[[1. 1. 1.]
 [1. 1. 1.]]

Forgive me if I am missing something very obvious.

[cgarciae]

Hey @deadsoul44, it works if you have more than 2 dimension e.g:

import jax.numpy as jnp

x = jnp.ones((8, 28, 28, 1))
y = x.reshape((x.shape[0], -1))
print(y.shape) # (8, 784)

else you the shape won't change as you note but all this is correct behaviour.

[cgarciae]

Note, flattening here means "flatten sample".

[deadsoul44]

Thanks a lot for the explanation.