Replace diagonal elements with vector

@SimonW Hi Simon, it works fine Sorry for that.

But will this operation affect the autograd for the vector and for the matrix A ?