def sddmm_cache_Y(op, req_grad_X, req_grad_Y):
    """Rules to identify whether to cache Y in SDDMM forward stage."""
    if op in ["mul", "dot"] and req_grad_X:
        return True
    return False
Rules to identify whether to cache Y in SDDMM forward stage.
python/dgl/backend/pytorch/sparse.py