Python – How to create a metric in Keras that returns multidimensional values?

I’m using keras to solve multiple types of problems. My data is very unbalanced, so I’m trying to create something similar to a confusion matrix. My dataset is very large and saved as HDF5, so I use HDF5Matrix to get X and Y, making the scikit-learn confusion matrix irrelevant (as far as I know).
I’ve seen that it’s possible save the predictions and true labels, or output the error per label, but a more elegant solution is to create a multidimensional quantity to accumulate (predicted, real) pairs of labels (sort of like a confusion matrix).
I use the following callback to try to see what each batch/epoch looks like:

from keras.callbacks import LambdaCallback
batch_print_callback = LambdaCallback(on_batch_end=lambda batch, logs: 
print(logs),on_epoch_end=lambda epoch, logs: print(logs))

But it only accumulates one value (usually the average of sorts).

I

also tried to see if I could return y_pred/y_true as follows (try to see if I can print multidimensional values in the log):

def pred(y_true, y_pred):
     return y_pred

def true(y_true, y_pred):
    return y_true

However, it doesn’t return multidimensional values as I expected
So basically, my question is, can I use keras to accumulate multidimensional metrics?

Well, as far as I know, this is not possible because K.mean is applied before returning the tensor value. I posted a issue about this on Keras GitHub.
The best design I came up with was a measure of each cell in the confusion matrix, and a callback to collect them, inspired by the thread mentioned in the question.
A workable solution can be found here