flax

Do not call `jnp` on shapes in AutoEncoder example

#1516

Merged