Theano自定义操作：定义grad（）

Question

我想创建一个自定义Op以便在PyMC3中使用。 该Op查找函数的根：f(x) = x + env*exp(x) - a*b^2，其中env = np.array([1, 2])，因此，寻根函数应为a和b的任何给定对返回向量。 这是我的定义方式：

from scipy import optimize
import numpy as np
import theano
import theano.tensor as tt

envod = np.array([1, 2])

def func(x, a, b, env):
    value = x + env * np.exp(x) - a * b**2
    return value

def jac(x, a, b, env):
    jac = 1 + env * np.exp(x)
    return jac

def x_from_ab(a, b, env):
    Len = len(env)
    value = np.zeros(Len)
    for i in range(len(envod)):
        value[i] = optimize.newton(func, 1, fprime = jac, args = (a, b, env[i]))
    return value

class Xf(tt.Op):
    itypes = [tt.dscalar, tt.dscalar]
    otypes = [tt.dvector]

    def perform(self, node, inputs, outputs):
        a, b = inputs
        x = x_from_ab(a, b, envod)
        outputs[0][0] = np.array(x)

    def grad(self, inputs, output_gradients):
        a, b = inputs
        x = self(a, b)
        g, = output_gradients
        return [-g[0] * (-b**2)/(1 + envod[0] * tt.exp(x[0])), -g[0] * (-2*a*b)/(1 + envod[0] * tt.exp(x[0]))]

我想知道应该如何定义grad()？我已经阅读了所有找到的帖子/文档。欢迎任何建议或指向有用参考的链接。

非常感谢！