Example: Writing a simple neural network# Note This is a placeholder for a section in the new JAX tutorials.