a_selected = tf$gather(a, a_index)
b_selected = tf$gather(b, b_index)
c_selected = tf$gather(c, c_index)

logit = a_selected * b_selected - c_selected
y_model = tf$sigmoid(logit)

loss = tf$losses$log_loss(y_real, y_model)

a_mm = mm = tf$nn$moments(a, axes = 0L)

total_loss = loss +
    tf$abs(a_mm[[1]]) +   # mean(a) = 0
    tf$abs(a_mm[[2]] - 1) # var(a) = 1^2

a_sort = tf$contrib$framework$sort(a)
n = length(a)
ideal_dist = qnorm((1:n)/(n+1), mean = 0, sd = 1)
divergence = tf$reduce_mean(tf$abs(a_sort - ideal_dist))

total_loss = loss +
    divergence