Is there a way to get feature specific information while calculating loss function on a batch in Keras

I'm trying to build a ANN regression model that regress to a bulk property of the associated data.

For example, assume I have a dataset with 3 columns A, B ,C and a target. based on the target, the probability density of the population for column A will be different.

I want to create a custom loss function to regress to these PDFs within the same batch (i.e, not generate a-priori distributions but depend on the batch population)

If I were to do it by hand it would look something like this:

def LossWrapper(input_batch_data,y_pred,y_truth):
 
    def loss(y_pred,y_truth): 

        PDF_truth = GeneratePDF(inputBatchData.A,y_truth)     

        PDF_pred  = GeneratePDF(inputBatchData.A,y_pred)          

        squared_difference = tf.square(PDF_truth - PDF_pred  )          
        return    tf.reduce_mean(squared_difference, axis=-1)      
    return loss(y_pred,y_truth) 

But I haven't found a way to access the input data.

Any suggestions?