I created a py_function to wrap the numpy interp function. But when I take the gradient, it returns a single number, when it should be a vector of length 10. What's going on? I don't know know to debug since py_functions are under the hood.
import tensorflow as tf
import numpy as np
# Define the Python function that uses numpy's interp
def numpy_interp(x, xp, fp):
return np.interp(x, xp, fp)
# Create a TensorFlow function using tf.py_function
def tf_interp(x, xp, fp):
return tf.py_function(numpy_interp, [x, xp, fp], tf.float32)
## now for the test:
input = tf.Variable(tf.random.uniform(shape=(10,),dtype=tf.float32))
vals = tf.random.uniform(shape=(10,),dtype=tf.float32)
with tf.GradientTape() as tape:
outputs = tf_interp(vals, input, tf.range(10,dtype=tf.float32))# [x, xp, yp]
print('outputs=',outputs)
print('outputs shape=',outputs.shape)
dy_dx = tape.gradient(outputs, input)
print(dy_dx)
If I replace the output of numpy_interp with x+xp+fp instead of np.interp(x, xp, fp), it gives the correct shape. If I replace it with np.sin(x+xp+fp) it gives the same problem. So I think it's a problem with calling numpy functions.
