Love to see it. A perfect example of why this optimization can't be done automatically - in the case of `else` you're working with a mutable reference to `x` passed in, which means that now your function is mutating something it used to not mutate.
A "safe" way to do this is still straightforward, I think.
from copy import copy
def _forward(self, x, context=None):
x = x.contiguous() if x.device.type == 'mps' else x
x = copy(x)
x += self.attn1(self.norm1(x))
x += self.attn2(self.norm2(x), context=context)
x += self.ff(self.norm3(x))
return x
It could be faster but I don't know what `x` is and I'm not going to guess. Also, `copy` may not be sufficient, `deepcopy` may be necessary - again, I don't know what `x` is so I can't figure that out. Pls use type annotations :)
How about this? (As the copy operation implicit in x=x+y seemed ok.)
def _forward(self, x, context=None):
x = x.contiguous() if x.device.type == 'mps' else x
x = x + self.attn1(self.norm1(x))
x += self.attn2(self.norm2(x), context=context)
x += self.ff(self.norm3(x))
return x
staticassertion's point is that the current code's usage of `+=` mutates the x that was passed in by the caller, and their suggestion is to copy x into a function local before mutating it, which is similar to how the original `+` code also worked on a function local x (the result of `attn1() + x`).
That's not the problem though. The problem is that the += operations mutate x in place, but the right hand side reads from x. There is no copy you could insert like that to fix this. You would have to do the following, for example
Instead of
x = op(x) + x -> x += op(x)
Do
x_copy = copy(x)
x += op(x_copy)
If you do
x_copy += op(x_copy)
Then you are still mutating x_copy while op() reads it.
EDIT:
I also don't think copy(x) will copy the actual tensor data, although I'm not super familiar with Pytorch.
Comments
Love to see it. A perfect example of why this optimization can't be done automatically - in the case of `else` you're working with a mutable reference to `x` passed in, which means that now your function is mutating something it used to not mutate.
A "safe" way to do this is still straightforward, I think.
It could be faster but I don't know what `x` is and I'm not going to guess. Also, `copy` may not be sufficient, `deepcopy` may be necessary - again, I don't know what `x` is so I can't figure that out. Pls use type annotations :)How about this? (As the copy operation implicit in x=x+y seemed ok.)
Alternatively, keep the first line as is. That gives you a copy that’s only known to the function, so you can change the later ones.
I would only do that if I had seen it to be faster, though, and add a comment on why the first line couldn’t do +=.
That's not safe if the problem is the in place mutation. You will still mutate x while reading from it.
staticassertion's point is that the current code's usage of `+=` mutates the x that was passed in by the caller, and their suggestion is to copy x into a function local before mutating it, which is similar to how the original `+` code also worked on a function local x (the result of `attn1() + x`).
That's not the problem though. The problem is that the += operations mutate x in place, but the right hand side reads from x. There is no copy you could insert like that to fix this. You would have to do the following, for example
Instead of x = op(x) + x -> x += op(x)
Do
x_copy = copy(x) x += op(x_copy)
If you do x_copy += op(x_copy)
Then you are still mutating x_copy while op() reads it.
EDIT: I also don't think copy(x) will copy the actual tensor data, although I'm not super familiar with Pytorch.
Ah, ok, I assumed the issue was happening outside of the function. If the issue is actually those intermediaries being mutated, bummer.
As for copy vs deepcopy, like I said, I have no idea what the type is so I don't know that deepcopy is necessary or not.
In x += op(x)
mutation of LHS x really starts before RHS has been evaluated completely?
I assume so, although if I remember correctly the actual error in the thread is caused by the the gradient computation
Might be platform dependent whether that first line counts as a mutate or not, seeing as it can be converted to not do anything in some cases.
All of the lines that have += mutate x in place
Unless x is immutable, in which case the semantics of both variants agree.