<p dir="auto">I tried getting the inverse (related to issue <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="437943844" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/646" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/646/hovercard" href="https://github.com/google/jax/issues/646">#646</a>), and it worked fine sort of. Here is some pseudocode of what I was trying to do.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def get_distances(...):
  svd code here to get inverse matrix
  return distances

def energy(params, positions, cell):
   get distances, do calculation
   return energy

def forces(params, positions, cell):
    dedr = grad(energy, 1)
    return -dedr(...)  # works fine

def stress(params, position, cell):
   def e0(params, positions, cell, strain):
       modify positions and cell with strain
       return energy(params, newpositions, newcell)

   deds = grad(e0, 3)
   return deds(params, positions, cell, np.zeros(3)) "><pre class="notranslate"><code class="notranslate">def get_distances(...):
  svd code here to get inverse matrix
  return distances

def energy(params, positions, cell):
   get distances, do calculation
   return energy

def forces(params, positions, cell):
    dedr = grad(energy, 1)
    return -dedr(...)  # works fine

def stress(params, position, cell):
   def e0(params, positions, cell, strain):
       modify positions and cell with strain
       return energy(params, newpositions, newcell)

   deds = grad(e0, 3)
   return deds(params, positions, cell, np.zeros(3)) 
</code></pre></div>
<p dir="auto">The last function gave me this error:<br>
<code class="notranslate">NotImplementedError: Forward-mode differentiation rule for 'svd' not implemented</code></p>
<p dir="auto">I guess the issue is that the strain modifies the cell, which is the array I need an inverse of, and this triggers the need for a derivative through the svd, and it is evidently not implemented.</p>
<p dir="auto">It seems like a clear error, so this is mostly a report I guess that it is not implemented. It isn't obvious where that would be done. Any hints?</p>