Heaviside (unit step) activation in TensorFlow

Question

I need to implement a perceptron in TensorFlow, however, the heaviside (unit step) activation seems not available in TensorFlow. It is not in tf., not in tf.nn., not in tf.keras.activations.. I guess because TensorFlow is gradient-based library and heaviside activation has no gradient.

I wonder why this basic function is not there. Any work-around for this? to make a perceptron.

Answer 1

TensorFlow has no heaviside (unit step) activation function possibly because TF is gradient-based library and heaviside has no gradient. I had to implement my own heaviside using the decorator @tf.custom_gradient:

#Heaviside (Unit Step) function with grad
@tf.custom_gradient
def heaviside(X):
  List = [];

  for I in range(BSIZE):
    Item = tf.cond(X[I]<0, lambda: tf.constant([0], tf.float32), 
                           lambda: tf.constant([1], tf.float32));  
    List.append(Item);

  U = tf.stack(List);

  #Heaviside half-maximum formula
  #U = (tf.sign(X)+1)/2;

  #Div is differentiation intermediate value
  def grad(Div):
    return Div*1; #Heaviside has no gradient, use 1.

  return U,grad;