Fix introductory CNN example in recent JAX.

Fix introductory CNN example in recent JAX. #1015