Defining grad() for custom Theano Op that solves nonlinear system of equations

Hi all,

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[0], F[1], F[2] ]

def solve_system(parameters, constants):
   # x0: some initial guess for roots
   roots = scipy.optimize.root(system, x0, args = (parameters, constants))
   return roots.x[2]  # 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[0]
         vector_inputs = inputs[1]

         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[0][0] = 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!