Class RotatE

include/instance/model/knowledge_graph.h:446–575

 */
template<class _Vector>
class RotatE {
public:
    static_assert(_Vector::dim % 2 == 0, "Model `RotatE` can only be instantiated with even-dimensional vectors");
    static const size_t dim = _Vector::dim;
    typedef _Vector Vector;
    typedef typename _Vector::Float Float;

    __host__ __device__
    static void forward(const Vector &head, const Vector &tail, const Vector &relation, Float &output, float margin) {
        output = 0;
        FOR(i, dim / 2) {
            Float h_re = head[i * 2];
            Float h_im = head[i * 2 + 1];
            Float t_re = tail[i * 2];
            Float t_im = tail[i * 2 + 1];
            Float phase = relation[i];
            Float r_re = cos(phase);
            Float r_im = sin(phase);
            Float distance_re = h_re * r_re - h_im * r_im - t_re;
            Float distance_im = h_re * r_im + h_im * r_re - t_im;
            output += sqrt(distance_re * distance_re + distance_im * distance_im);
        }
        output = margin - SUM(output);
    }

    template<OptimizerType optimizer_type>
    __host__ __device__
    static void backward(Vector &head, Vector &tail, Vector &relation,
                         float margin, Float gradient, const Optimizer &optimizer,
                         float relation_lr_multiplier = 1, Float weight = 1) {
        auto update = get_update_function<Float, optimizer_type>();
        FOR(i, dim / 2) {
            Float phase = relation[i];
            Float r_re = cos(phase);
            Float r_im = sin(phase);
            Float h_re = head[i * 2];
            Float h_im = head[i * 2 + 1];
            Float t_re = tail[i * 2];
            Float t_im = tail[i * 2 + 1];
            Float distance_re = h_re * r_re - h_im * r_im - t_re;
            Float distance_im = h_re * r_im + h_im * r_re - t_im;
            Float grad = gradient / (sqrt(distance_re * distance_re + distance_im * distance_im) + kEpsilon);
            // head
            Float head_re_grad = -grad * (distance_re * r_re + distance_im * r_im);
            Float head_im_grad = -grad * (-distance_re * r_im + distance_im * r_re);
            head[i * 2] -= (optimizer.*update)(h_re, head_re_grad, weight);
            head[i * 2 + 1] -= (optimizer.*update)(h_im, head_im_grad, weight);
            // tail
            tail[i * 2] -= (optimizer.*update)(t_re, grad * distance_re, weight);
            tail[i * 2 + 1] -= (optimizer.*update)(t_im, grad * distance_im, weight);
            // relation
            Float relation_grad =
                    -grad * (distance_re * (h_re * -r_im + h_im * -r_re) + distance_im * (h_re * r_re + h_im * -r_im));
            relation[i] -= relation_lr_multiplier * (optimizer.*update)(phase, relation_grad, weight);
        }
    }

    template<OptimizerType optimizer_type>