Repository: pykeen/pykeen
Issue: Extended Partition/Backfill Representation
Root cause: Unimplemented Feature Gap

### Problem Statement

I'm working on some biomedical examples and have a need for a representation like the Backfill, but not quite.

As it is, the backfill takes a subset of entities with known representations and puts them in a static embedding, then makes the rest trainable.

What I want to do is create a backfill that also has a trainable representation in addition to the static ones I give


```mermaid
graph LR
   s(Static)
   l(Learnable - Known)
   u(Learnable - Unknown) 
   t1(Learnable Transformation)
   s --> t1 --> c[Combination] --> p[Partition]
   l --> c
   u --> p
```

Where:

- `n` entities, $n_k$ have static representations, $n_u$ do not
- the $d_k$-dimensional static representations for known entities $S \in \mathbb{R}^{n_k \times d_k}$
- (optional) transformation $T \in \mathbb{R}^{d_k \times t}$. This can also be no-op if the combination takes care of this
- the $d_u$-dimensional learnable weights for known entities $W_1 \in \mathbb{R}^{n_u \times d_u}$


Because there are many different kinds of [combinations](https://pykeen.readthedocs.io/en/stable/api/pykeen.nn.combination.Combination.html#pykeen.nn.combination.Combination) already implemented in PyKEEN, the results from T1 and T2 don't have to be the same dimensional, but this might make things easier. T1 and T2 might also be possible to skip since the combination itself could handle this part, e.g., the [ConcatCombination](https://pykeen.readthedocs.io/en/stable/api/pykeen.nn.combination.ConcatCombination.html) stacks two tensors, so they can be different lengths

### Describe the solution you'd like

1. We need a tutorial for [pykeen.nn.CombinedRepresentation](https://pykeen.readthedocs.io/en/stable/api/pykeen.nn.representation.CombinedRepresentation.html#pykeen.nn.representation.CombinedRepresentation)
2. Let's make another macro-level representation for accomplishing what I'm trying to do, which is one step more complicated than just the backfill representation since I want to have multiple entity types with pre-defined vectors


```mermaid
graph LR
   s1(Chemical Fingerprints)
   l1(Learnable Chemical Embeddings)
   t1(Learnable Chemical Transformation)
   s2(Protein Fingerprints)
   l2(Learnable Protein Embeddings)
   t3(Learnable Protein Transformation)
   u(Learnable - Unknown) 
   
   s1 --> t1 --> c1[Chemical Combination] --> p[Partition]
   l1 --> c1
   s2 --> t3 --> c2[Protein Combination] --> p[Partition]
   l2 --> c2
   u --> p
```

Even better, I might be adding multiple different static vectors. Chemicals can have structural fingerprints, high-throughput fingerprints, etc.

### Describe alternatives you've considered

It's already possible to achieve a smaller version of this with the BackfillRepresentation

### Additional information

_No response_

### Issue Template Checks

- [x] This is not a bug report (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)

Sounds like an interesting problem, and a very nice use case for showing how to combine different representations from PyKEEN.

I quickly drafted something which may correspond to the described setup (minus the actual biomedical information 😄 ).

```python
"""Example for biochemical use case."""

import torch
from torch import nn

from pykeen.nn import CombinedRepresentation, Embedding, Representation, TransformedRepresentation, init
from pykeen.nn.representation import PartitionRepresentation


def _create_combined(features: torch.Tensor, output_dim: int = 32, hidden_dim: int = 32) -> Representation:
    max_id, dim = tuple(features.shape)
    # create chemical representations
    repr_features = Embedding(
        max_id=max_id, embedding_dim=dim, trainable=False, initializer=init.PretrainedInitializer(tensor=features)
    )
    transformed_features = TransformedRepresentation(
        transformation=nn.Sequential(nn.Linear(dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, output_dim)),
        base=repr_features,
    )
    embedding = Embedding(max_id=max_id, embedding_dim=hidden_dim)
    # TODO: max_id=None should be default
    return CombinedRepresentation(max_id=None, base=(transformed_features, embedding))


def _create_multi_backfill(max_id: int, bases: list[Representation], ids: list[list[int]]) -> Representation:
    assignment = torch.zeros(max_id, 2, dtype=torch.long)
    shape: tuple[int, ...] | None = None
    seen: set[int] = set()
    for base_id, (this_ids, base) in enumerate(zip(ids, bases, strict=True), start=1):
        # shape validation
        if shape is None:
            shape = base.shape
        elif shape != tuple(base.shape):
            raise ValueError(f"Shape mismatch {shape} vs. {base.shape}")
        # index overlap validation
        if duplicate_ids := seen.intersection(this_ids):
            raise ValueError(f"Encountered {duplicate_ids=}")
        seen.update(this_ids)
        # update assignment
        for local_id, global_id in enumerate(this_ids):
            assignment[global_id, 0] = base_id
            assignment[global_id, 1] = local_id
    assert shape is not None
    # build backfill
    missing = sum(len(this_ids) for this_ids in ids)
    if missing < 0:
        raise ValueError("max_id is too small.")
    backfill = Embedding(max_id=max_id - missing, shape=shape)
    # combine
    return PartitionRepresentation(assignment=assignment, bases=[backfill, *bases])


max_id = 32
chem_ids = [1, 3, 7]
protein_ids = [2, 8, 12, 17]

# create chemical representations
chem_feat = torch.rand(len(chem_ids), 32)
repr_chem = _create_combined(features=chem_feat)

# create protein representations, features may have different shape
protein_feat = torch.rand(len(protein_ids), 24)
repr_protein = _create_combined(features=protein_feat)

# embedding for everything else
repr = _create_multi_backfill(max_id=max_id, bases=[repr_chem, repr_protein], ids=[chem_ids, protein_ids])
```

`_create_multi_backfill` sounds like something that is generic enough to make it into its own component (or be included in `BackfillRepresentation`). `_create_combined` might be too specific to the desired transformation though.