I'm trying to use TensorFlow to minimize the below loss function (L) with respect to u. There are 3 variables, u, x_opt, L, with the following dependency graph:
u ---(f)--> x_opt ---(g)--> L,
with the exact form of the dependency governed by functions f and g.
def f(u):
def f_helper(u,x):
# with u held fixed, f_helper is a convex function of x
# the exact form of f_helper does not matter
return np.linalg.norm(x-u)
curried_f_helper = lambda x: f_helper(u,x)
x_opt = scipy.optimize(curried_f_helper,np.random.uniform(5))['x']
return x_opt
def g(x_opt):
# the exact form of g does not matter
return np.ones(x_opt.shape).dot(x_opt)
def L(u):
# want to optimize L over u
x_opt = f(u)
return g(x_opt)
# use TensorFlow to minimize L over u...
The complication is that f() does not have an analytical functional form - u parameterizes an optimization problem whose solution is x_opt. So TensorFlow would not be able to compute the gradient of f with respect to u. However, I can use implicit differentiation to manually compute this gradient. Ideally, I'd be able to define a new op representing f, and register its gradient (that I manually calculate).
My question is: How should I implement the op representing f and specify its gradient? Is it possible to define the op for f using only Python, and if so, will I have to use tf.pyfunc?
