jitting through wrapper

[SaitejaUtpala]

let cost be real-valued function takes two arguments, call it params, data.

def cost(params, data):
    # does some computation 
    return cost_value

now computing gradient wrt params would amount

gradient = (grad(cost))(params, data)

Now, assume that I have to build a new type of gradient called new_type_grad, where after computing a gradient of functioncost, there is transformation called transform parametrized by again params(first argument of cost) that transforms grad.So,

new_type_gradient = transform(  (grad(cost))(params,data),     params)

so I need to develop a higher order function new_type_grad that would automaically do as above, i.e.,

new_type_gradient =(new_type_grad(cost))(params, data)

Further I should be able to jit the new_type_grad( just we would jit the grad )

new_type_gradient = (jit(new_type_grad(cost)))(params, data)

Here is the first approach I thought

def new_type_grad(cost):
    def inner(params, theta):
        new_type_gradient = transform( grad(cost))(params,data), params)
        return new_type_gradient
    return inner

key problem is that (jit(new_type_grad))(cost) is significantly slower than just new_type_grad(cost). Any better approach?

Thanks in advance.

[jakevdp]

Sorry for the slow response here – I find it hard to understand what you are asking. Can you provide some actual minimal code and timings that show what slowdown you're seeing?