Problem Statement

I have a TensorFlow machine learning model that should learn to mimic a digital result [0, 1], the output of the model is a float number, so I have used a MeanSquaredError and planned to rescale the output of the model to [0, 1] by adding the absolute minimum value of the output and then dividing by the maximum value of the result, and then round the result of the last operation to ensure that the inputs given to the MeanSquaredError is digital.

Code Sample

import tensorflow as tf

mse = tf.keras.losses.MeanSquaredError()

def loss(y_true, y_pred):  
  yp = y_pred + tf.abs(tf.reduce_min(y_pred))
  yp = yp / tf.reduce_max(yp)
  yp = tf.round(yp)
  
  return mse(y_true, yp)

Attempts

I tried to look onto the error and it mentioned that the loss function must be differentiable and tried the solutions provided in other questions which was concentrating on tf.round but nothing worked.

One of them was this

x_rounded_NOT_differentiable = tf.round(x)
x_rounded_differentiable = x - tf.stop_gradient(x - x_rounded_NOT_differentiable)

it seems that tf.reduce_min and tf.reduce_max are non-differentiable too.

Remarks

I hope that my presentation was clear, and I am a bit new to TensorFlow but not to programming, so is there a better way to write this?

Edits

I have recently found a solution but still need the advice.
the solution was to multiply then clip the value by 0 and 1 like:

import tensorflow as tf

mse = tf.keras.losses.MeanSquaredError()

def loss(y_true, y_pred):
  yp = tf.clip_by_value(y_pred*100, 0, 1)
  return mse(y_true, yp)