Python optimize if tree

I have "if decision tree" and I want to know if it's possible to optimize it:

    def s(a, b):
    """
    :param a: 0,1,2
    :param b: 0,1,2
    :return: 0,1,2
    """
    if a == 0:
        if b == 0:
            return 2
        elif b == 1:
            return 2
        else:  # b == 2
            return 0
    elif a == 1:
        if b == 0:
            return 2
        elif b == 1:
            return 2
        else:  # b == 2
            return 1
    else:  # a==2
        if b == 0:
            return 0
        elif b == 1:
            return 1
        else:  # b==2
            return 2

All cases included, I'm trying use (a,b) == (1,2) or (a,b) == (2,1) return 1 and the same when return 0, other cases return 2, but it is slower.