Skip to content

Commit

Permalink
assert out per-channel on "quantized tensor: add preliminary support …
Browse files Browse the repository at this point in the history
…for advanced indexing, try 2"

Summary:

This is less ambitious redo of
#49129.

We make the

```
xq_slice = xq[:, [0], :, :]
```

indexing syntax work if `xq` is a quantized Tensor.  For now, we are
making the code not crash, with an in efficient `dq -> index -> q`
implementation.  A future PR can optimize performance by removing
the unnecessary memory copies (which will require some non-trivial
changes to TensorIterator).

Test Plan:

```
python test/test_quantization.py TestQuantizedOps.test_advanced_indexing
```

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: [D25539365](https://our.internmc.facebook.com/intern/diff/D25539365)

[ghstack-poisoned]
  • Loading branch information
vkuzo committed Dec 14, 2020
1 parent 230cbc2 commit 0f959c4
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions aten/src/ATen/native/TensorAdvancedIndexing.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,11 @@ Tensor index(const Tensor & self, TensorList indices) {
}

Tensor quantized_index(const Tensor & self, TensorList indices) {
TORCH_INTERNAL_ASSERT(
self.qscheme() == c10::kPerTensorAffine ||
self.qscheme() == c10::kPerTensorSymmetric,
"Indexing for per-channel quantized Tensors is not supported.");

// For now, this is a naive implementation which does dq -> index -> q.
// TODO(future PR): improve performance by removing the copies.
const auto& self_dq = self.dequantize();
Expand Down

0 comments on commit 0f959c4

Please sign in to comment.