I’m writing a custom Theano Op that wraps
scipy.optimize.root() to find the roots of a nonlinear system of three equations and three unknowns, though I’m only interested in the value of the third root. In pseudocode, it looks like:
def system(unknowns, parameters, constants): return [ F, F, F ] def solve_system(parameters, constants): # x0: some initial guess for roots roots = scipy.optimize.root(system, x0, args = (parameters, constants)) return roots.x # return the value of the third root
I’ve also written a Theano Op that successfully returns the root of interest as a vector, since I’m looping over a vector of constants and want to see this root at the various constant values. In pseudocode again:
class MyCustomOp(theano.Op): itypes = [ tt.dscalar, tt.dvector ] otypes = [ tt.dvector ] def __init__(self, constants): self.constants = constants def perform(self, node, inputs, outputs): scalar_input = inputs vector_inputs = inputs vector_of_roots = np.empty(N) for i in range(N): vector_of_roots[i] = solve_system(parameters = [scalar_inputs, vector_inputs], constants = self.constants) outputs = vector_of_roots
Now this works perfectly fine, I’ve tested it out and
MyCustomOp returns an array of correct expected values. My question is: How would I define the grad() so that I can use this in NUTS sampling?
I’ve got two inputs, one scalar (let’s call it S) and one vector (V). There’s three equations in the system. There’s N values in
vector_of_roots. As I understand,
grad() needs to return a list of size 2, but I don’t know how to format this list.
Thanks in advance!