CODEWITHSUNDEEP
tensorflow
munichmath
munichmath

Reputation: 75

Tensorflow custom activation function with tf.cond

I'm trying to write a custom activation function using tf.custom_gradient. Specifically I want to use the taylor expansion of 1/x for x<1 and 1/x otherwise. Here's my code:

@tf.custom_gradient
def taylor_inverse(x):
    def func(x):
        return(tf.cond(x<1, taylor(x), tf.math.reciprocal(x)))

    def grad(upstream):
        return(tf.cond(upstream<1, taylor_grad(upstream), inv_diff(upstream)))
    return func(x), grad

@tf.function
def taylor(x):
    return(4 - 6 * x + 4 * x ** 2 - x ** 3)
    
@tf.function
def taylor_grad(x):
    return(-3 * x ** 2 + 8 * x - 6)

@tf.function
def inv_diff(x):
    return(-tf.math.reciprocal(x)**2)

I get the error message:

TypeError: 'Tensor' object is not callable

Equations are -x3+4x2-6x+4 and for the gradient -3x2+8x-6, and I get error in this line:

layer_inverse = Lambda(lambda x: taylor_inverse(x),output_shape=(1,))(layer)

Thank you for your help

Upvotes: 2

Views: 146

Answers (1)

Kaveh
Kaveh

Reputation: 4960

tf.cond second and third arguments should be callable function. So, use it like this:

@tf.custom_gradient
def taylor_inverse(x):
    def func(x):
        return(tf.cond(x<1, lambda: taylor(x), lambda: tf.math.reciprocal(x)))

    def grad(upstream):
        return(tf.cond(upstream<1, lambda: taylor_grad(upstream), lambda: inv_diff(upstream)))
    return func(x), grad

Upvotes: 1

Related Questions