<p dir="auto">Hi, the following attempt to take the gradient of a function involving a matrix exponential</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
import jax
from jax.scipy.linalg import expm

def foo(mat):
    return jnp.sum(expm(mat))

grad = jax.grad(foo)(jnp.eye(2))"><pre class="notranslate"><code class="notranslate">import jax.numpy as jnp
import jax
from jax.scipy.linalg import expm

def foo(mat):
    return jnp.sum(expm(mat))

grad = jax.grad(foo)(jnp.eye(2))
</code></pre></div>
<p dir="auto">is failing with the following assertion error</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="~/anaconda3/envs/envname/lib/python3.7/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     96   pval_primals, pval_tangents = tree_unflatten(out_tree(), out_pvals)
     97   aval_primals, const_primals = unzip2(pval_primals)
---&gt; 98   assert all(aval_primal is None for aval_primal in aval_primals)
     99   if not has_aux:
    100     return const_primals, pval_tangents, jaxpr, consts

AssertionError: "><pre class="notranslate"><code class="notranslate">~/anaconda3/envs/envname/lib/python3.7/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     96   pval_primals, pval_tangents = tree_unflatten(out_tree(), out_pvals)
     97   aval_primals, const_primals = unzip2(pval_primals)
---&gt; 98   assert all(aval_primal is None for aval_primal in aval_primals)
     99   if not has_aux:
    100     return const_primals, pval_tangents, jaxpr, consts

AssertionError: 
</code></pre></div>
<p dir="auto">I'm using <code class="notranslate">jax.__version__  == 0.1.62</code> and <code class="notranslate">jaxlib.__version__ == 0.1.43</code>, thanks!</p>
<p dir="auto">The pull request <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="554529148" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2062" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/2062/hovercard" href="https://github.com/google/jax/pull/2062">#2062</a> implements the jvp of the matrix exponential by using the Fréchet derivative. I think the vjp can be easily added by transposing the input matrix in <code class="notranslate">expm_frechet</code>:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="defvjp(expm, lambda g, ans, matrix: expm_frechet(_T(matrix), g, compute_expm=False))"><pre class="notranslate"><span class="pl-en">defvjp</span>(<span class="pl-s1">expm</span>, <span class="pl-k">lambda</span> <span class="pl-s1">g</span>, <span class="pl-s1">ans</span>, <span class="pl-s1">matrix</span>: <span class="pl-en">expm_frechet</span>(<span class="pl-en">_T</span>(<span class="pl-s1">matrix</span>), <span class="pl-s1">g</span>, <span class="pl-s1">compute_expm</span><span class="pl-c1">=</span><span class="pl-c1">False</span>))</pre></div>
<p dir="auto">I can test and implement this solution if <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="554529148" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2062" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/2062/hovercard" href="https://github.com/google/jax/pull/2062">#2062</a> is merged.</p>