Repository: pykeen/pykeen
Issue: predict.consumer_scores requires a square tensor of entities and relations in PartiallyRestrictedPredictionDataset
Root cause: Algorithmic Error

### Describe the bug

**How I encountered it**
_I think this a bug and not a limitation. My apologies if I am mistaken._
I encountered the problem while having a large list of tails on a certain relation, which I wanted to collect the top k predicted heads over all tails in my list. 

**The bug:**
When using consume_scores on a PartiallyRestrictedPredictionDataset while not containing the same amount of entities and/or relations in the dataset(i.e, not a square tensor/matrix), if yields for instance an error like this: 

`IndexError: index 2 is out of bounds for dimension 0 with size 2`

This occurs everytime we reach the end of the biggest inner square matrix of the two given restrictions in the dataset.

The error occurs on line 923 in predict.py, when creating a batch and it visits __getitem__ on line 878. 
If we have 3 heads and 3 relations, we get 9 batches successfully, and no errors. However, if they are not equal, like 3x2, it ends up with the error above.

----



### How to reproduce

To use the same code snippet as commented in predict.PartialRestrictedPredictionDataset, except we remove an entity ("uk") from the heads, making the self.parts in __getitem__ line 878, none-square (2 x 3), expecting it to calculate for 6 combinations. The error occurs on the third iteration in this example:

```python
from pykeen.predict import PartiallyRestrictedPredictionDataset
heads = result.training.entities_to_ids(entities=["netherlands", "poland"])
relations = result.training.relations_to_ids(relations=["reltourism", "tourism", "tourism3"])
dataset = PartiallyRestrictedPredictionDataset(heads=heads, relations=relations)

# calculate all scores for this restricted set, and keep k=3 largest
from pykeen.predict import consume_scores, TopKScoreConsumer
consumer = TopKScoreConsumer(k=3)
consume_scores(result.model, ds, consumer)
score_pack = consumer.finalize()

# add labels
df = result.training.tensor_to_df(score_pack.result, score=score_pack.scores)
```


### Environment

| Key             | Value                    |
|-----------------|--------------------------|
| OS              | posix                    |
| Platform        | Linux                    |
| Release         | 5.15.0-52-generic        |
| Time            | Sat Nov  5 15:49:47 2022 |
| Python          | 3.9.7                    |
| PyKEEN          | 1.9.1-dev                |
| PyKEEN Hash     | UNHASHED                 |
| PyKEEN Branch   |                          |
| PyTorch         | 1.11.0                   |
| CUDA Available? | true                     |
| CUDA Version    | 11.3                     |
| cuDNN Version   | 8200                     |

### Additional information

_No response_

### Issue Template Checks

- [X] This is not a feature request (use a different issue template if it is)
- [X] This is not a question (use the [discussions](https://github.com/pykeen/pykeen/discussions/new) forum instead)
- [X] I've read the text explaining why including environment information is important and understand if I omit this information that my issue will be dismissed

Hi @AImenes ,

> The error occurs on line 923 in predict.py, when creating a batch and it visits getitem on line 878.

Would you please provide a full traceback?

Sure!

**Code**:
```python
# train model; note: needs larger number of epochs to do something useful ;-)
from pykeen.pipeline import pipeline
result = pipeline(dataset="nations", model="mure", training_kwargs=dict(num_epochs=0))

# create prediction dataset, where the head entities is from a set of European countries,
# and the relations are connected to tourism
from pykeen.predict import PartiallyRestrictedPredictionDataset
tail = result.training.entities_to_ids(entities=["netherlands", "poland"])
relations = result.training.relations_to_ids(relations=["reltourism", "tourism", "tourism3"])
dataset = PartiallyRestrictedPredictionDataset(tails=tail, relations=relations, target="head")

# calculate all scores for this restricted set, and keep k=3 largest
from pykeen.predict import consume_scores, TopKScoreConsumer
consumer = TopKScoreConsumer(k=3)
consume_scores(result.model, dataset, consumer)
score_pack = consumer.finalize()

# add labels
df = result.training.tensor_to_df(score_pack.result, score=score_pack.scores)
print(df)
```
**Traceback**:
```python-traceback
No random seed is specified. Setting to 1513198918.
Training epochs on cuda:0: 0epoch [00:00, ?epoch/s]
INFO:pykeen.evaluation.evaluator:Starting batch_size search for evaluation now...
INFO:pykeen.evaluation.evaluator:Concluded batch_size search with batch_size=201.
Evaluating on cuda:0: 100%|████████████████████████████████████| 201/201 [00:00<00:00, 15.8ktriple/s]
INFO:pykeen.evaluation.evaluator:Evaluation took 0.02s seconds
scoring:   0%|                                                        | 0.00/6.00 [00:00<?, ?batch/s]
iteration: (relation:0, entity: 0) (user-added print statement)
iteration: (relation:0, entity: 1) (user-added print statement)
iteration: (relation:0, entity: 2) (user-added print statement) #This entity doesnt exist
Traceback (most recent call last):                                                                   
  File "/Coding/Embedding/pykeen_test.py", line 15, in <module>
    consume_scores(result.model, dataset, consumer)
  File "/home/anaconda3/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/anaconda3/lib/python3.9/site-packages/torch_max_mem/api.py", line 293, in inner
    result, self.parameter_value[h] = wrapped(*args, **kwargs)
  File "/home/anaconda3/lib/python3.9/site-packages/torch_max_mem/api.py", line 193, in wrapper_maximize_memory_utilization
    func(*bound_arguments.args, **p_kwargs, **bound_arguments.kwargs),
  File "/home/anaconda3/lib/python3.9/site-packages/pykeen/predict.py", line 924, in consume_scores
    for batch in tqdm(data_loader, desc="scoring", unit="batch", unit_scale=True, leave=False):
  File "/home/anaconda3/lib/python3.9/site-packages/tqdm/std.py", line 1180, in __iter__
    for obj in iterable:
  File "/home/anaconda3/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 530, in __next__
    data = self._next_data()
  File "/home/anaconda3/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 570, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/anaconda3/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/anaconda3/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 49, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/anaconda3/lib/python3.9/site-packages/pykeen/predict.py", line 879, in __getitem__
    return torch.as_tensor([self.parts[0][quotient], self.parts[1][remainder]])
IndexError: index 2 is out of bounds for dimension 0 with size 2
```

The fix: 

line 878:
`return torch.as_tensor([self.parts[0][quotient], self.parts[1][remainder]])
`
change to
line 878:
`return torch.as_tensor([self.parts[0][remainder], self.parts[1][quotient]])
`


Re-opening until this is auto-closed by #1157 being merged into `master`.