add nominal support for int32 indices in index/index_put ops #86309
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/86309
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 FailuresAs of commit 3a243e0: The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
lezcano
left a comment
lezcano
left a comment
There was a problem hiding this comment.
This will be quite slow. Can this path be accessed from the Python API? If so, this may bite people.
| TORCH_CHECK_INDEX(false, "tensors used as indices must be long, byte or bool tensors"); | ||
| if (allow_int) { | ||
| if (scalarType != kLong && scalarType != kByte && scalarType != kBool && scalarType != kInt) { | ||
| TORCH_CHECK_INDEX(false, "tensors used as indices must be long, byte or bool tensors"); |
|
This path is accessible from python, but it's not slower than it is today for users who have to use int32 indices for whatever reason (they'd need to manually convert them to int64), so I don't consider it "biting". When used in the decomposition the backend will generate a kernel directly for int32 indices. |
|
Yes... but also no. If you are creating an integer tensor, let it be for indexing or whatever, by default you are going to create it with dtype |
|
Integer tensors by default are created with int64 dtype (functions like |
|
@pytorchbot merge -f "test failures unrelated" |
|
@pytorchbot successfully started a merge job. Check the current status here. |
|
Hey @ngimel. |
…#86309) Summary: Currently index_select/index_add decompositions decompose to `index` or `index_put` ops. The problem with this is that `index_select` and `index_add` accept int32 indices while `index` doesn't. That leads to error in meta func for those decompositions. This PR adds non-performant support for int32 indices to `index` operations, to allow decompositions go through. Pull Request resolved: #86309 Approved by: https://github.com/lezcano Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/dc9c507d24d0c833cb09105177326f1f6bbe99c4 Reviewed By: seemethere Differential Revision: D40166997 Pulled By: seemethere fbshipit-source-id: 812a8b8e5d7088ba20586ced63ea294ea994fa0f
Currently index_select/index_add decompositions decompose to
indexorindex_putops. The problem with this is thatindex_selectandindex_addaccept int32 indices whileindexdoesn't. That leads to error in meta func for those decompositions. This PR adds non-performant support for int32 indices toindexoperations, to allow decompositions go through.