Repository: pykeen/pykeen
Issue: Use torch.nn.functional.embedding instead of integer indexing
GitHub labels: enhancement

**Is your feature request related to a problem? Please describe.**
https://github.com/pytorch/pytorch/issues/15245 suggests to use `torch.nn.functional.embedding` instead of integer indexing to accelerate the backward pass.

**Describe the solution you'd like**
Use this instead of integer indexing (e.g. when looking up representations).

**Describe alternatives you've considered**
It is not clear whether the performance advantage will persist, and the code might loose some of its readability.

**Additional context**
PyTorch Geometric switched to using it, cf. https://github.com/pyg-team/pytorch_geometric/pull/97

fascinating. is this because using that call optimizes it?

I am not exactly sure why - maybe because integer indexing allows more than embedding, and thus needs to have some extra code paths.

Reading the full linked PyTorch thread, it seems that the performance difference has much decreased since the initial post, and now is somewhat comparable (although still slower) on the benchmark example given there.

The PyG issue is also pretty old by now.

So either we would need to run some benchmarking for our use cases (looking at the transition to generic interaction functions ... :smile: ), or leave it for now.

I ran a small benchmarking script, but could not observe any coherent advantage. I suggest, we leave it as it is.

```python
import torch
from torch.utils.benchmark import Timer
from torch.utils.benchmark.utils.compare import Compare
Compare(
    results=[
        Timer(
            stmt=stmt,
            setup=f"x=torch.rand({n}, {d}, requires_grad=True, device=device); i=torch.randint({n}, size={s}, device=device)",
            description=f"n={n}, i={s}",
            env=f"d={device}",
            globals=dict(
                torch=torch,
                device=device,
            ),
            num_threads=num_threads,
        ).blocked_autorange()
        for n, d, s in [
            (15_000, 200, (256,)),
            (15_000, 200, (1024,)),
            (15_000, 200, (256, 256)),
        ]
        for num_threads in (1, 16)
        for stmt in ["x[i]", "torch.nn.functional.embedding(i, x)"]
        for device in [torch.device("cpu"), torch.device("cuda")]
    ],
).print()
```

#### RTX 2080 Ti
```
[---------------------------------------------------------  ---------------------------------------------------------]
                                                 |  n=15000, i=(256,)  |  n=15000, i=(1024,)  |  n=15000, i=(256, 256)
1 threads: -----------------------------------------------------------------------------------------------------------
  (d=cpu)   x[i]                                 |         15.1        |         48.7         |         22096.0       
            torch.nn.functional.embedding(i, x)  |         13.4        |         41.8         |         21896.8       
  (d=cuda)  x[i]                                 |         11.7        |         11.8         |           202.8       
            torch.nn.functional.embedding(i, x)  |          9.9        |         10.0         |           223.0       
16 threads: ----------------------------------------------------------------------------------------------------------
  (d=cpu)   x[i]                                 |         16.9        |         19.4         |          6247.3       
            torch.nn.functional.embedding(i, x)  |         18.3        |         19.2         |          6429.2       
  (d=cuda)  x[i]                                 |         12.0        |         11.8         |           202.8       
            torch.nn.functional.embedding(i, x)  |         10.0        |         10.0         |           223.3    

Times are in microseconds (us).
```

#### RTX 8000
```
[---------------------------------------------------------  ---------------------------------------------------------]
                                                 |  n=15000, i=(256,)  |  n=15000, i=(1024,)  |  n=15000, i=(256, 256)
1 threads: -----------------------------------------------------------------------------------------------------------
  (d=cpu)   x[i]                                 |         24.0        |        135.4         |         37048.8       
            torch.nn.functional.embedding(i, x)  |         16.6        |        118.7         |         38404.4       
  (d=cuda)  x[i]                                 |         15.3        |         15.2         |           194.1       
            torch.nn.functional.embedding(i, x)  |         13.5        |         13.6         |           211.5       
16 threads: ----------------------------------------------------------------------------------------------------------
  (d=cpu)   x[i]                                 |        189.6        |        300.9         |         42149.1       
            torch.nn.functional.embedding(i, x)  |        176.6        |        266.3         |         38892.1       
  (d=cuda)  x[i]                                 |         16.8        |         17.1         |           192.9       
            torch.nn.functional.embedding(i, x)  |         15.7        |         15.7         |           211.1       

Times are in microseconds (us).
```

Good this is as expected. I’m thinking that the embedding function is just a way to be more functional, but we don’t necessarily need that. Gonna close this