Adding support for hvd LaidOutVariable about mesh

[EiffL]

ok found a couple of problems, the update seems to work now, will create a proper branch ^^

[tobias-liaudat]

Hi @EiffL, I'm leaving this comment as a checkpoint of when we finished the IDRIS hackaton.

With @b-remy we were trying to solve this issue but we hadn't had the time to finish. I uploaded my progress in the branch tob_vars and Benjamin on his branch ben-variable.

With our current implementation we were able to run a forward pass of a simple dense network. However, when we tried to optimise the network we were having an error and we're not sure of the origin and how to solve it.

The error we found was this one, and is always related with the Assign and the optimisation.

Traceback (most recent call last):
  File "optim_demo.py", line 174, in <module>
    (type(obj).__name__, types_str))
TypeError: Can not convert a Assign into a Tensor or Operation.

TypeError: Fetch argument <mesh_tensorflow.ops.Assign object at 0x14d1fcb94c50> has invalid type <class 'mesh_tensorflow.ops.Assign'>, must be a string or Tensor. (Can not convert a Assign into a Tensor or Operation.)

To reproduce the error one can run this test script with this job script. The implementation of the LaidOutVariable is here.

[EiffL]

Thanks so much @tobias-liaudat will take a look!

[EiffL]

oookkkkk so I tried something here:
https://github.com/DifferentiableUniverseInitiative/mesh/tree/u/EiffL/toy_model

with this script https://github.com/DifferentiableUniverseInitiative/mesh/blob/u/EiffL/toy_model/examples/toy_model_gpu.sh

It runs apparently, can save and restore, but not clear if it's actually training ^^" the loss function doesnt go down much

[EiffL]

The goal here would be to be able to train the MNIST demo model with the new backend implemenation https://github.com/DifferentiableUniverseInitiative/mesh/blob/master/examples/mnist.py

[EiffL]

I've opened branch variables starting from tob_vars, and cleaned it up a bit.
Problems where:

• lowering needs to happen after computing the mesh gradients and stuff, otheriwse the update ops are not registered in the lowering
• small problem in the assignement of slices

[EiffL]

probably, we can try to first train the mnist model with commenting out the restore and save parts