98e43878ec0e56f4bacd2960169a483af8f44c31,app/ops.py,,pit_mse_loss,#Any#Any#Any#Any#Any#,262

Before Change

        s_loss

    """
    ndim = len(s_x.get_shape().as_list())
    assert -ndim <= pit_axis < ndim
    pit_axis %= ndim
    assert pit_axis != 0
    reduce_axes = [

After Change

            tf.square(s_x - s_y), reduce_axes)
        s_loss_sets = tf.einsum(
            "bij,pij->bp", s_cross_loss, s_perms_onehot)
        s_loss_sets_idx = tf.argmin(s_loss_sets, axis=1)
        s_loss = tf.gather_nd(
            s_loss_sets,
            tf.stack([
                tf.range(hparams.BATCH_SIZE, dtype=tf.int64),
                s_loss_sets_idx], axis=1))
        s_loss = tf.reduce_mean(s_loss)
    return s_loss, v_perms, s_loss_sets_idx

Instances