Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add OpInfo for torch.Tensor.__setitem__ (to test prims.copy_with_setitem) #1357

Open
kshitij12345 opened this issue Oct 28, 2024 · 0 comments
Labels

Comments

@kshitij12345
Copy link
Collaborator

kshitij12345 commented Oct 28, 2024

Currently we have a manual test at

def test_setitem():
def fn(a):
a[:3] = 2
return a * 2
a_ref = torch.ones(5)
out_ref = fn(a_ref)
a = torch.ones(5)
jf = thunder.jit(fn)
out = jf(a)
assert_close(a, a_ref)
assert_close(out, out_ref)

and PR #1322 extends this test.

I think it would be nice if we can leverage the OpInfo infra though I am not sure how it interacts for inplace operations (as currently all ops in opinfo are out-of-place).

We have some infra which uses OpInfo to generate a new list _inplace_opinfo for inplace ops by finding their corresponding inplace version. (Maybe we can directly add the OpInfo for setitem here?)

@ops(_inplace_opinfos, supported_dtypes=(dtypes.float32,))
def test_functionalization(op: OpInfo, device: str, dtype: dtypes.dtype, executor, _):

cc @apaz-cli

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

No branches or pull requests

1 participant