Keras custom gradient function for discontinuous functions

Question

I tried writing a custom layer, which should take a predictions from a model, test the prediction and return the result together with the original input to the next layer. M