Кусочное ограничение в Pytorch

Question

Допустим, я предсказываю кватернион. Как реализовать это кусочное ограничение, которое обусловлено скалярным / вещественным «w» (для кватерниона = [x, y, z, w]).

In Python:

def positive_real_quat(quat):
    """
    Flip quaternion to positive real hemisphere so there is a single quaternion describing each rotation

    :param quat: [x,y,z,w]
    :return:
    """
    x, y, z, w =quat
    if w >= 0:
        return quat
    else:
        return -quat

Но мне нужно, чтобы вышеприведенный код был дифференцируемым как конечная функция в моей нейронной сети. Как мне это сделать?