Creating custom gradient for `tf.abs` may solve the problem:


    @tf.custom_gradient
    def abs_with_grad(x):
      y = tf.abs(x);
    
      def grad(div): # Derivation intermediate value
        g = 1; # Use 1 to make the chain rule just skip abs
        return div*g;
    
      return y,grad;