# Buggy code
def split(
    self,
    ratios: Union[float, Sequence[float]] = 0.8,
    *,
    random_state: RandomHint = None,
    randomize_cleanup: bool = False,
) -> List["TriplesFactory"]:
    """Split a triples factory into a train/test.

    :param ratios: There are three options for this argument. First, a float can be given between 0 and 1.0,
     non-inclusive. The first triples factory will get this ratio and the second will get the rest. Second,
     a list of ratios can be given for which factory in which order should get what ratios as in ``[0.8, 0.1]``.
     The final ratio can be omitted because that can be calculated. Third, all ratios can be explicitly set in
     order such as in ``[0.8, 0.1, 0.1]`` where the sum of all ratios is 1.0.
    :param random_state: The random state used to shuffle and split the triples in this factory.
    :param randomize_cleanup: If true, uses the non-deterministic method for moving triples to the training set.
     This has the advantage that it doesn't necessarily have to move all of them, but it might be slower.

    .. code-block:: python

        ratio = 0.8  # makes a [0.8, 0.2] split
        training_factory, testing_factory = factory.split(ratio)

        ratios = [0.8, 0.1]  # makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)

        ratios = [0.8, 0.1, 0.1]  # also makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)
    """
    n_triples = self.triples.shape[0]

    # Prepare shuffle index
    idx = np.arange(n_triples)
    random_state = ensure_random_state(random_state)
    random_state.shuffle(idx)

    # Prepare split index
    if isinstance(ratios, float):
        ratios = [ratios]

    ratio_sum = sum(ratios)
    if ratio_sum == 1.0:
        ratios = ratios[:-1]  # vsplit doesn't take the final number into account.
    elif ratio_sum > 1.0:
        raise ValueError(f"ratios sum to more than 1.0: {ratios} (sum={ratio_sum})")

    sizes = [int(split_ratio * n_triples) for split_ratio in ratios]
    # Take cumulative sum so the get separated properly
    split_idxs = np.cumsum(sizes)

    # Split triples
    triples_groups = np.vsplit(self.triples[idx], split_idxs)
    logger.info(
        "done splitting triples to groups of sizes %s",
        [triples.shape[0] for triples in triples_groups],
    )

    # Make sure that the first element has all the right stuff in it
    logger.debug("cleaning up groups")
    triples_groups = _tf_cleanup_all(
        triples_groups, random_state=random_state if randomize_cleanup else None
    )
    logger.debug("done cleaning up groups")

    for i, (triples, exp_size, exp_ratio) in enumerate(
        zip(triples_groups, sizes, ratios)
    ):
        actual_size = triples.shape[0]
        actual_ratio = actual_size / exp_size * exp_ratio
        if actual_size != exp_size:
            logger.warning(
                f"Requested ratio[{i}]={exp_ratio:.3f} (equal to size {exp_size}), but got {actual_ratio:.3f} "
                f"(equal to size {actual_size}) to ensure that all entities/relations occur in train.",
            )

    # Make new triples factories for each group
    return [
        TriplesFactory.from_labeled_triples(
            triples=triples,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            compact_id=False,
        )
        for triples in triples_groups
    ]


# Traceback (KeyError)
Traceback (most recent call last):
File "/Users/cthoyt/dev/pykeen/scratch/tst.py", line 49, in <module>
main()
File "/Users/cthoyt/dev/pykeen/scratch/tst.py", line 36, in main
testing.new_with_restriction(relations=evaluation_relation_whitelist)
File "/Users/cthoyt/dev/pykeen/src/pykeen/triples/triples_factory.py", line 641, in new_with_restriction
relations = list(relations) + list(map(self.relation_to_inverse.__getitem__, relations))
KeyError: 'accusation_inverse'

# Fixed code
def split(
    self,
    ratios: Union[float, Sequence[float]] = 0.8,
    *,
    random_state: TorchRandomHint = None,
    randomize_cleanup: bool = False,
) -> List["TriplesFactory"]:
    """Split a triples factory into a train/test.

    :param ratios: There are three options for this argument. First, a float can be given between 0 and 1.0,
     non-inclusive. The first triples factory will get this ratio and the second will get the rest. Second,
     a list of ratios can be given for which factory in which order should get what ratios as in ``[0.8, 0.1]``.
     The final ratio can be omitted because that can be calculated. Third, all ratios can be explicitly set in
     order such as in ``[0.8, 0.1, 0.1]`` where the sum of all ratios is 1.0.
    :param random_state: The random state used to shuffle and split the triples in this factory.
    :param randomize_cleanup: If true, uses the non-deterministic method for moving triples to the training set.
     This has the advantage that it doesn't necessarily have to move all of them, but it might be slower.

    .. code-block:: python

        ratio = 0.8  # makes a [0.8, 0.2] split
        training_factory, testing_factory = factory.split(ratio)

        ratios = [0.8, 0.1]  # makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)

        ratios = [0.8, 0.1, 0.1]  # also makes a [0.8, 0.1, 0.1] split
        training_factory, testing_factory, validation_factory = factory.split(ratios)
    """
    # input normalization
    ratios = normalize_ratios(ratios)
    generator = ensure_torch_random_state(random_state)

    # convert to absolute sizes
    sizes = get_absolute_split_sizes(n_total=self.num_triples, ratios=ratios)

    # Split indices
    idx = torch.randperm(self.num_triples, generator=generator)
    idx_groups = idx.split(split_size=sizes, dim=0)

    # Split triples
    triples_groups = [self.mapped_triples[idx] for idx in idx_groups]
    logger.info(
        "done splitting triples to groups of sizes %s",
        [triples.shape[0] for triples in triples_groups],
    )

    # Make sure that the first element has all the right stuff in it
    logger.debug("cleaning up groups")
    triples_groups = _tf_cleanup_all(
        triples_groups, random_state=generator if randomize_cleanup else None
    )
    logger.debug("done cleaning up groups")

    for i, (triples, exp_size, exp_ratio) in enumerate(
        zip(triples_groups, sizes, ratios)
    ):
        actual_size = triples.shape[0]
        actual_ratio = actual_size / exp_size * exp_ratio
        if actual_size != exp_size:
            logger.warning(
                f"Requested ratio[{i}]={exp_ratio:.3f} (equal to size {exp_size}), but got {actual_ratio:.3f} "
                f"(equal to size {actual_size}) to ensure that all entities/relations occur in train.",
            )

    # Make new triples factories for each group
    return [
        self.clone_and_exchange_triples(mapped_triples=triples)
        for triples in triples_groups
    ]