from tinyml.core import Backend as np
from tinyml.layers import Softmax

def cross_entropy_with_softmax_loss(predicted, ground_truth):
    softmax = Softmax('softmax')
    output_probabilities = softmax(predicted)
    print(output_probabilities)
    loss = np.mean(-np.log(output_probabilities[np.arange(output_probabilities.shape[0],dtype=np.int8), ground_truth]+1e-20))
    
    output_probabilities[np.arange(output_probabilities.shape[0]), ground_truth] -= 1
    gradient = output_probabilities / predicted.shape[0]
    return loss, gradient