title: Combine scipy.root and Jax Jacobian
body: <p>I am having trouble using the Jacobian from JAX with <code>scipy.root</code>. In the below example, the <code>root</code> works without the Jacobian, while it fails with the Jacobian. Any ideas on what I need to rewrite in order to get the code below working with the Jacobian?</p>
  <pre><code>from jax import jacfwd
  from scipy.optimize import root
  import numpy as np
  
  def objectFunction(valuesEndo, varNamesEndo, valuesExo, varNamesExo, equations): 
    for i in range(len(varNamesExo)):
        exec(&quot;%s = %.10f&quot; %(varNamesExo[i], valuesExo[i]))
  
    for i in range(len(varNamesEndo)):
      exec(&quot;%s = %.10f&quot; %(varNamesEndo[i], valuesEndo[i]))
      
    equationVector = np.zeros(len(equations))
    for i in range(len(equations)):
        exec('equationVector[%d] = eval(equations[%d])' %(i, i))    
        
    return equationVector
  
  varNamesEndo = ['x', 'y']
  valuesEndoInitialGuess = [1., 1.]
  
  varNamesExo = ['a', 'b']
  valuesExo = [1., 1.]
  
  equations = ['a*x+b*y**2-4',
              'np.exp(x) + x*y - 3']
  
  equations = ['a*x**2 + b*y**2',
              'a*x**2 - b*y**2']
  
  # Without Jacobian
  sol1 =  root(fun=objectFunction,
              x0=valuesEndoInitialGuess, 
              args=(varNamesEndo, valuesExo, varNamesExo, equations))
  #----&gt; Works
  
  # With Jacobian
  jac  = jacfwd(objectFunction)
  sol2 =  root(fun=objectFunction,
              x0=valuesEndoInitialGuess, 
              args=(varNamesEndo, valuesExo, varNamesExo, equations),
              jac=jac)
  #----&gt; Not woring
  </code></pre>
  <p>At least there seems to be problems with the line</p>
  <pre><code>for i in range(len(varNamesEndo)):
          exec(&quot;%s = %.10f&quot; %(varNamesEndo[i], valuesEndo[i]))
  </code></pre>