How to define the grad_func_paramter for a function for defvjp ?

[BhattacharyaSumit]

Hi,
I am a newbie with autograd. I am having difficulty in writing a proper gradient function to use in the defvjp function.

@autograd.primitive
def compute_Q(U , xs, phi, args):
  edof, penalS = args['edof'] , args['penalS']
  j0, jl, beta = args['j0'], args['jl'], args['beta']
  array = beta * anp.dot(U[edof],phi)
  one = scipy.special.expit(array)
  two = scipy.special.expit(-array)
  three = anp.divide(one, two+ 1e-10  , dtype = anp.float64)
  je = jl*xs**penalS + j0 * (1/4) * anp.sum(three-1,1)
  je_check = je < -jl
  args['je_check'] = je_check
  #print(je_check)
  je[je_check]=0
  Q = anp.zeros((args['nelx']+1)*(args['nely']+1)).ravel()

  for i in range(edof.shape[0]):
    for j in range(edof[i].shape[0]):
      Q[edof[i][j]] =Q[edof[i][j]] + je[i] * (1/4) * args['ela']
  return Q

autograd.extend.defvjp(compute_Q , something, something, something... )
Here phi and args are constant so their grad_func should be None, in my opinion.
U and xs are 1d vectors and edof is a 2D vector of indices in a particular order.
Could you help me in structuring the grad_fun for U and xs?

def grad_compute_Q_U(ans , U, xs, phi, args):
  def jvp(v):
    g = lambda U : compute_Q(U ,  xs,  phi, args)
    return autograd.jacobian(g)(U) @ v 
  return jvp 

Is this the proper way to write the grad function?

[j-towns]

Hi there! Can you say what kind of error you're getting? I think the problem might be that you're using autograd.jacobian from within a vjp definition — this might lead to a recursion depth error because jacobian uses vjp (see here). Two more questions:

1. Does compute_Q need to be a primitive? We use primitives when we need to define a derivative of a function which is not differentiable using Autograd, but compute_Q looks like it should be, what happens if you remove the primitive decorator and try to do autograd.grad(compute_Q)(U, xs, phi, args)?

2. If that doesn't work, maybe you do need to wrap compute_Q in a primitive, or maybe you want to define some custom/efficient derivative by hand. In that case you can't rely on autograd.jacobian or other autograd operators in the vjp definition, you have to write it manually.

[j-towns]

Make sure to do

import autograd.scipy as scipy

[BhattacharyaSumit]

Hi @j-towns!
I tried it without using primitives, but autograd is giving an error stating that je[je_check] = 0 is not possible, since autograd does not allow assignment. So, I just used

je = je._value
je[je_check] = 0

But, in that case autograd no longer traces je as it has been converted to to je._value.

Owing to all this, I wrapped the function with primitive decorator.

[BhattacharyaSumit]

Make sure to do

import autograd.scipy as scipy

And, in my code I need to use scipy.sparse which is not supported by autograd.scipy, so I think using autograd.scipy may not work for me

[BhattacharyaSumit]

When I don't use primitive, this error also pops up

TypeError: float() argument must be a string or a number, not 'ArrayBox'