Implement torch.tril_indices and torch.triu_indices (#12653) (#14904)
authorShen Li <shenli@fb.com>
Wed, 12 Dec 2018 23:18:57 +0000 (15:18 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Wed, 12 Dec 2018 23:40:14 +0000 (15:40 -0800)
commit90f9e8103c9b1108ae923f5ab67bcc566e942699
tree5d8ec0eb90518785f743ab9b1502f879187ea8a6
parent342e62f1e3849118034958cc012d68b1b5e25418
Implement torch.tril_indices and torch.triu_indices (#12653) (#14904)

Summary:
This is an optimized implementation that does the following:

1. created an empty Tensor of correct size.
2. fill the Tensor with correct values.

The following three designs to fill in the Tensor result in roughly the same performance. Hence, the 2nd option is taken for simpler code, and to return contiguous tensors.

1. Sequential: fill row coordinates first, then columns. This results in two for-loop and more arithmetic operations.
2. Interleaved: fill in index coordinates one by one, which jumps between the two output Tensor rows in every iteration.
3. Transpose: create a n X 2 Tensor, fill the Tensor sequentially, and then transpose it.

<img width="352" alt="screen shot 2018-12-10 at 3 54 39 pm" src="https://user-images.githubusercontent.com/16999635/49769172-07bd3580-fc94-11e8-8164-41839185e9f9.png">

NOTE:

This implementation returns a 2D tensor, instead of a tuple of two tensors. It means that users will not be able to do the following:

```python
x = torch.ones(3, 3)
i = torch.tril_indices(3, 3)
x[i]  # need to first convert the 2D tensor into a tuple of two 1D tensors.
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14904

Reviewed By: zou3519

Differential Revision: D13433027

Pulled By: mrshenli

fbshipit-source-id: 41c876aafcf584832d7069f7c5929ffb59e0ae6a
aten/src/ATen/native/TensorFactories.cpp
aten/src/ATen/native/native_functions.yaml
test/test_cuda.py
test/test_torch.py
tools/autograd/gen_python_functions.py
torch/_torch_docs.py