Lightning-AI / Lightning-AI/lightning-thunder

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

Open
#1,357 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

testing
Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

Currently we have a manual test at https://github.com/Lightning-AI/lightning-thunder/blob/1575797a8d21e0b845e188ddd6af305be6457b1c/thunder/tests/test_ops.py#L242-L253

and PR https://github.com/Lightning-AI/lightning-thunder/pull/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?)

https://github.com/Lightning-AI/lightning-thunder/blob/1575797a8d21e0b845e188ddd6af305be6457b1c/thunder/tests/test_inplace_functionalization.py#L108-L109

cc @mruberry @apaz-cli

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the manual test in thunder/tests/test_ops.py around lines 242-253 and review PR #1322, then inspect the _inplace_opinfo handling in thunder/tests/test_inplace_functionalization.py around lines 108-109. Determine how torch.Tensor.__setitem__ should be represented in OpInfo for prims.copy_with_setitem; done means the OpInfo-based coverage replaces or extends the manual coverage and the relevant tests pass.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
testing-qa
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
32/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.