def METHOD_NAME(self, stage: str) -> List[Tuple[str, str]]:
    """
    Returns the pairs of input and output states that the given stage
    operates upon.
    """
    out = [
        (s, e) for (s, e, st) in self.graph.edges(data="stage") if st.name == stage
    ]
    assert len(out) > 0, f"No state tranformation for {stage} found."
    return out