from typing import Any


def recursive_state_dict(x: Any) -> Any:
    """
    Recursively call state_dict() on all elements contained in a list / tuple /
    dict so that it can be saved safely via torch.save().

    Args:
        x: any nesting of list / tuple / dict containing state_dict()able
            objects

    Returns:
        the same structure with state dicts
    """
    if hasattr(x, 'state_dict'):
        return x.state_dict()
    if isinstance(x, tuple):
        return tuple(recursive_state_dict(xx) for xx in x)
    if isinstance(x, list):
        return [recursive_state_dict(xx) for xx in x]
    if isinstance(x, dict):
        return {k: recursive_state_dict(v) for k, v in x.items()}
    return x