On-device Parameter Operations

I would like to implement Averaged SGD in dynet as done in https://github.com/pytorch/pytorch/blob/master/torch/optim/asgd.py

The issue is implementing https://github.com/pytorch/pytorch/blob/master/torch/optim/asgd.py#L75:

`state['ax'].add_(p.data.sub(state['ax']).mul(state['mu']))`

`state['ax']` and `p.data` are an on-GPU Tensors, `add_` is an in-place operation.

My implementation in dynet is using numpy arrays:

`state['ax'] += (p.as_array() - state['ax']) * state['mu']`

`state['ax']` is an np-array and `p` is a Parameter.

This works but is VERY slow. I would like to do this operation on the GPU using as little memory as possible (in-place `add_`).

If I define `state['ax']` as a Parameter, I think I can do the subtraction and multiplication using the computational graph, but how would I do the in-place addition?

Thanks!

Start with torch/optim/asgd.py, especially the linked update at line 75, then trace DyNet's Parameter handling and the current NumPy-based implementation described in the issue. Done means the averaged-SGD update runs on the GPU, uses an in-place operation with minimal extra memory, and preserves the intended update behavior.

Written by the indexing model from the issue text.