I am trying to make a custom layer in Tensorflow. It is supposed to be an activation function with trainable parameters. I am supposed to apply a function to all the element