Repository navigation
Fix grad copies - #1854
Fix grad copies#1854
Conversation
angeloskath
left a comment
angeloskath
left a comment
There was a problem hiding this comment.
Weird that this went like this for so long 🤔.
| py_value_out = fun(*tree[0], **tree[1]); | ||
|
|
||
| tree_replace(tree, arrays, a); | ||
| tree_fill(tree, arrays); |
There was a problem hiding this comment.
I would expect that not doing this would result in the *args and **kwargs containing a which can be placeholders. But it doesn't happen, do we need the tree_fill here?
There was a problem hiding this comment.
But it doesn't happen
You tested it and it's not the case? That is strange and surprising if so..
There was a problem hiding this comment.
I think I did. I ran the test on this PR on mlx before this PR. The list arrays still contains the original inputs. Nothing gets broken 🤷♂️
There was a problem hiding this comment.
tree_replace(tree, arrays, a) shouldn't do anything as it replaces with the placeholders again.
There was a problem hiding this comment.
So the tests definitely pass without it.. but nevertheless I think we should keep it. Here's an example of some odd behavior if we remove it:
import mlx.core as mx
def fun(arrs):
return arrs[0]
arrs = [mx.array(1.0)]
init_id = id(arrs[0])
mx.grad(fun)(arrs)
final_id = id(arrs[0])
print(init_id, final_id) # ids won't matchThere was a problem hiding this comment.
In fact, maybe I should add that as a test :)
There was a problem hiding this comment.
I think it makes sense that it doesn't break. We don't actually use placeholders a la vmap. We just set them to tracers. They keep their primitives and everything. I don't know why I expected them to break. When not tracing you can evaluate tracers fine.
Totally agree that the tree_fill should be there.
53bda9c to
65131a2
Compare
* fix grad with copies * add test * add test
2 participants
Closes #1853