import cv2
import numpy as np
import torch
import torch.nn.functional as F
import tensorflow as tf

y_true=np.random.randn(1,2,2)
y_pred=np.random.randn(1,2,2)
print('y_true',y_true, '\t y_pred', y_pred)

def tf_kld (y_true, y_pred):
    
    y_true = tf.convert_to_tensor(y_true)
    y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.reshape(y_true, [1,-1])
    y_pred = tf.reshape(y_pred, [1,-1])
    y_true=tf.nn.softmax(y_true)
    y_pred=tf.nn.softmax(y_pred)
    print('tf softmax y_pred:', y_pred)
    kld_loss = tf.reduce_sum(y_true*tf.math.log(y_true/y_pred), axis=1)
    return kld_loss

def torch_kld(y_true, y_pred):
    y_true = torch.from_numpy(y_true)
    y_pred = torch.from_numpy(y_pred)
    y_true = y_true.view(1, -1)
    y_pred = y_pred.view(1, -1)
    y_true = F.softmax(y_true, dim=1)
    y_pred = F.log_softmax(y_pred, dim=1)
    print('torch softmax y_pred:', y_pred)
    # print('torch kld loss:', (y_true*torch.log(y_true/y_pred)).sum(-1))
    kld = F.kl_div(y_pred, y_true, reduction='none')
    kld_loss = kld.sum(-1)
    return kld_loss

print('tensorflow kld loss:', tf_kld(y_true, y_pred))
print('pytorch kld loss:', torch_kld(y_true, y_pred))