Adding support for LaidOutVariables

[EiffL]

@tobias-liaudat @b-remy I'm opening this PR to work on a common solution to #1

This script seems to be working.... at least something gets updated.

What remains to be done:

• [ ] : add support for restore hook
• [ ] : try to train mnist model