<p dir="auto">Would it be possible to implement this interesting paper idea with jax / flax?</p>
<p dir="auto">Intelligent Matrix Exponentiation<br>
paper: <a href="https://arxiv.org/pdf/2008.03936.pdf" rel="nofollow">https://arxiv.org/pdf/2008.03936.pdf</a><br>
code: <a href="https://github.com/google-research/google-research/tree/master/m_layer">https://github.com/google-research/google-research/tree/master/m_layer</a><br>
wiki: <a href="https://en.wikipedia.org/wiki/Matrix_exponential" rel="nofollow">https://en.wikipedia.org/wiki/Matrix_exponential</a></p>
<p dir="auto">expm docs: <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.expm.html#jax.scipy.linalg.expm" rel="nofollow">https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.expm.html#jax.scipy.linalg.expm</a></p>
<p dir="auto">expm frechet looks like the right thing to make the gradient:<br>
<a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.expm_frechet.html#jax.scipy.linalg.expm_frechet" rel="nofollow">https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.expm_frechet.html#jax.scipy.linalg.expm_frechet</a></p>
<p dir="auto">I tried a few options, here's a simple one:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax
from flax import nn
jnp = jax.numpy
vec_expm = jnp.vectorize(jax.scipy.linalg.expm, signature='(k)-&gt;(k)')

@nn.module
def MLayer(x, D=D_CODE):
    x = nn.Dense(x, D ** 2)
    x = x.reshape(x.shape[:-1] + (D, D))
    x = vec_expm(x
    x = x.reshape(x.shape[:-2] + (D ** 2,))
    x = nn.Dense(x, D)
    return x"><pre class="notranslate"><code class="notranslate">import jax
from flax import nn
jnp = jax.numpy
vec_expm = jnp.vectorize(jax.scipy.linalg.expm, signature='(k)-&gt;(k)')

@nn.module
def MLayer(x, D=D_CODE):
    x = nn.Dense(x, D ** 2)
    x = x.reshape(x.shape[:-1] + (D, D))
    x = vec_expm(x
    x = x.reshape(x.shape[:-2] + (D ** 2,))
    x = nn.Dense(x, D)
    return x
</code></pre></div>
<p dir="auto">however this crashes because of this:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead."><pre class="notranslate"><code class="notranslate">ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead.
</code></pre></div>
<p dir="auto">Would it work if we wire up expm_frechet primitive?</p>
<p dir="auto">a simpler reproduction without other nn stuff:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
import jax

rng = jax.random.PRNGKey(0)
x = jax.random.uniform(rng, (2, 2))

def f(x):
    y = jax.scipy.linalg.expm(x)
    credit = jnp.sum(jnp.abs(y))
    return credit, y

f = jax.value_and_grad(f)

c, y = f(x)
print(c, y)"><pre class="notranslate"><code class="notranslate">import jax.numpy as jnp
import jax

rng = jax.random.PRNGKey(0)
x = jax.random.uniform(rng, (2, 2))

def f(x):
    y = jax.scipy.linalg.expm(x)
    credit = jnp.sum(jnp.abs(y))
    return credit, y

f = jax.value_and_grad(f)

c, y = f(x)
print(c, y)
</code></pre></div>
<p dir="auto">full trace:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
&lt;ipython-input-14-f2d123ee1505&gt; in &lt;module&gt;
----&gt; 1 f(x)

~/miniconda3/lib/python3.8/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    485     tree_map(partial(_check_input_dtype_grad, holomorphic), dyn_args)
    486     if not has_aux:
--&gt; 487       ans, vjp_py = _vjp(f_partial, *dyn_args)
    488     else:
    489       ans, vjp_py, aux = _vjp(f_partial, *dyn_args, has_aux=True)

~/miniconda3/lib/python3.8/site-packages/jax/api.py in _vjp(fun, *primals, **kwargs)
   1514   if not has_aux:
   1515     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-&gt; 1516     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1517     out_tree = out_tree()
   1518   else:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    108 def vjp(traceable, primals, has_aux=False):
    109   if not has_aux:
--&gt; 110     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    111   else:
    112     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     95   _, in_tree = tree_flatten(((primals, primals), {}))
     96   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
---&gt; 97   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
     98   out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)
     99   assert all(out_primal_pval.is_known() for out_primal_pval in out_primals_pvals)

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out, bottom, trace_type)
    421   with core.new_master(trace_type, bottom=bottom) as master:
    422     fun = trace_to_subjaxpr(fun, master, instantiate)
--&gt; 423     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    424     assert not env
    425     del master

~/miniconda3/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    148     gen = None
    149 
--&gt; 150     ans = self.f(*args, **dict(self.params, **kwargs))
    151     del args
    152     while stack:

&lt;ipython-input-10-995182967638&gt; in wrapped(*args, **kwargs)
      1 def value_and_jacobian(fun):
      2     def wrapped(*args, **kwargs):
----&gt; 3         return fun(*args, **kwargs), jax.jacfwd(fun)(*args, **kwargs)
      4     return wrapped

~/miniconda3/lib/python3.8/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    491     dtype = dtypes.result_type(ans)
    492     tree_map(partial(_check_output_dtype_grad, holomorphic), ans)
--&gt; 493     g = vjp_py(np.ones((), dtype=dtype))
    494     g = g[0] if isinstance(argnums, int) else g
    495     if not has_aux:

~/miniconda3/lib/python3.8/site-packages/jax/api.py in _vjp_pullback_wrapper(cotangent_dtypes, io_tree, fun, py_args)
   1458              &quot;match type of corresponding primal output ({})&quot;)
   1459       raise TypeError(msg.format(_dtype(a), dtype))
-&gt; 1460   ans = fun(*args)
   1461   return tree_unflatten(out_tree, ans)
   1462 

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in unbound_vjp(pvals, jaxpr, consts, *cts)
    115     cts = tuple(map(ignore_consts, cts, pvals))
    116     dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
--&gt; 117     arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
    118     return map(instantiate_zeros, arg_cts)
    119 

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
    200         cts_in_avals = [v.aval for v in eqn.outvars]
    201         call_jaxpr, params = core.extract_call_jaxpr(eqn.primitive, eqn.params)
--&gt; 202         cts_out = get_primitive_transpose(eqn.primitive)(
    203             params, call_jaxpr, invals, cts_in, cts_in_avals)
    204       else:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in call_transpose(primitive, params, call_jaxpr, args, ct, _)
    486     new_params = update_params(new_params, map(is_undefined_primal, args),
    487                                [type(x) is not Zero for x in ct])
--&gt; 488   out_flat = primitive.bind(fun, *all_args, **new_params)
    489   return tree_unflatten(out_tree(), out_flat)
    490 primitive_transposes[core.call_p] = partial(call_transpose, call_p)

~/miniconda3/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1132 
   1133   def bind(self, fun, *args, **params):
-&gt; 1134     return call_bind(self, fun, *args, **params)
   1135 
   1136   def process(self, trace, fun, tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1124   else:
   1125     tracers = map(top_trace.full_raise, args)
-&gt; 1126     outs = primitive.process(top_trace, fun, tracers, params)
   1127   return apply_todos(env_trace_todo(), map(full_lower, outs))
   1128 

~/miniconda3/lib/python3.8/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1135 
   1136   def process(self, trace, fun, tracers, params):
-&gt; 1137     return trace.process_call(self, fun, tracers, params)
   1138 
   1139   def post_process(self, trace, out_tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in process_call(self, call_primitive, f, tracers, params)
    273     update_params = call_param_updaters.get(call_primitive)
    274     new_params = update_params(params, nz_tangents) if update_params else params
--&gt; 275     result = call_primitive.bind(f_jvp, *primals, *nonzero_tangents, **new_params)
    276     primal_out, tangent_out = tree_unflatten(out_tree_def(), result)
    277     return [JVPTracer(self, p, t) for p, t in zip(primal_out, tangent_out)]

~/miniconda3/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1132 
   1133   def bind(self, fun, *args, **params):
-&gt; 1134     return call_bind(self, fun, *args, **params)
   1135 
   1136   def process(self, trace, fun, tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1124   else:
   1125     tracers = map(top_trace.full_raise, args)
-&gt; 1126     outs = primitive.process(top_trace, fun, tracers, params)
   1127   return apply_todos(env_trace_todo(), map(full_lower, outs))
   1128 

~/miniconda3/lib/python3.8/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1135 
   1136   def process(self, trace, fun, tracers, params):
-&gt; 1137     return trace.process_call(self, fun, tracers, params)
   1138 
   1139   def post_process(self, trace, out_tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in process_call(self, primitive, f, tracers, params)
    179                   else PartialVal.unknown(mapped_aval(pval[0]))
    180                   for pval, is_mapped in zip(in_pvals, params['mapped_invars'])]
--&gt; 181     jaxpr, out_pvals, consts, env_tracers = self.partial_eval(
    182         f, in_pvals, partial(primitive.bind, **params))
    183     if primitive.map_primitive:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in partial_eval(self, f, pvals, app)
    279     f = trace_to_subjaxpr(f, self.master, False)
    280     f, aux = partial_eval_wrapper(f, tuple(in_avals))
--&gt; 281     out_flat, (out_avals, jaxpr, env) = app(f, *in_consts), aux()
    282     out_consts, consts = split_list(out_flat, [len(out_flat)-len(jaxpr.constvars)])
    283     out_pvs = map(PartialVal, zip(out_avals, out_consts))

~/miniconda3/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1132 
   1133   def bind(self, fun, *args, **params):
-&gt; 1134     return call_bind(self, fun, *args, **params)
   1135 
   1136   def process(self, trace, fun, tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1121   if top_trace is None:
   1122     with new_sublevel():
-&gt; 1123       outs = primitive.impl(fun, *args, **params)
   1124   else:
   1125     tracers = map(top_trace.full_raise, args)

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
    524 
    525 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
--&gt; 526   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
    527                                *unsafe_map(arg_spec, args))
    528   try:

~/miniconda3/lib/python3.8/site-packages/jax/linear_util.py in memoized_fun(fun, *args)
    222       fun.populate_stores(stores)
    223     else:
--&gt; 224       ans = call(fun, *args)
    225       cache[key] = (ans, fun.stores)
    226     return ans

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    595   else:
    596     pvals: Sequence[pe.PartialVal] = [pe.PartialVal.unknown(aval) for aval in abstract_args]
--&gt; 597     jaxpr, pvals, consts = pe.trace_to_jaxpr(
    598         fun, pvals, instantiate=False, stage_out=True, bottom=True)
    599   map(prefetch, it.chain(consts, jaxpr_literals(jaxpr)))

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out, bottom, trace_type)
    421   with core.new_master(trace_type, bottom=bottom) as master:
    422     fun = trace_to_subjaxpr(fun, master, instantiate)
--&gt; 423     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    424     assert not env
    425     del master

~/miniconda3/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    148     gen = None
    149 
--&gt; 150     ans = self.f(*args, **dict(self.params, **kwargs))
    151     del args
    152     while stack:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
    203             params, call_jaxpr, invals, cts_in, cts_in_avals)
    204       else:
--&gt; 205         cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
    206                                                          **eqn.params)
    207     cts_out = [Zero(v.aval) for v in eqn.invars] if cts_out is Zero else cts_out

~/miniconda3/lib/python3.8/site-packages/jax/lax/lax_control_flow.py in _while_transpose_error(*_, **kwargs)
    536 
    537 def _while_transpose_error(*_, **kwargs):
--&gt; 538   raise ValueError(&quot;Reverse-mode differentiation does not work for &quot;
    539                    &quot;lax.while_loop or lax.fori_loop. &quot;
    540                    &quot;Try using lax.scan instead.&quot;)

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead."><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
&lt;ipython-input-14-f2d123ee1505&gt; in &lt;module&gt;
----&gt; 1 f(x)

~/miniconda3/lib/python3.8/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    485     tree_map(partial(_check_input_dtype_grad, holomorphic), dyn_args)
    486     if not has_aux:
--&gt; 487       ans, vjp_py = _vjp(f_partial, *dyn_args)
    488     else:
    489       ans, vjp_py, aux = _vjp(f_partial, *dyn_args, has_aux=True)

~/miniconda3/lib/python3.8/site-packages/jax/api.py in _vjp(fun, *primals, **kwargs)
   1514   if not has_aux:
   1515     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-&gt; 1516     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1517     out_tree = out_tree()
   1518   else:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    108 def vjp(traceable, primals, has_aux=False):
    109   if not has_aux:
--&gt; 110     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    111   else:
    112     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     95   _, in_tree = tree_flatten(((primals, primals), {}))
     96   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
---&gt; 97   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
     98   out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)
     99   assert all(out_primal_pval.is_known() for out_primal_pval in out_primals_pvals)

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out, bottom, trace_type)
    421   with core.new_master(trace_type, bottom=bottom) as master:
    422     fun = trace_to_subjaxpr(fun, master, instantiate)
--&gt; 423     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    424     assert not env
    425     del master

~/miniconda3/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    148     gen = None
    149 
--&gt; 150     ans = self.f(*args, **dict(self.params, **kwargs))
    151     del args
    152     while stack:

&lt;ipython-input-10-995182967638&gt; in wrapped(*args, **kwargs)
      1 def value_and_jacobian(fun):
      2     def wrapped(*args, **kwargs):
----&gt; 3         return fun(*args, **kwargs), jax.jacfwd(fun)(*args, **kwargs)
      4     return wrapped

~/miniconda3/lib/python3.8/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    491     dtype = dtypes.result_type(ans)
    492     tree_map(partial(_check_output_dtype_grad, holomorphic), ans)
--&gt; 493     g = vjp_py(np.ones((), dtype=dtype))
    494     g = g[0] if isinstance(argnums, int) else g
    495     if not has_aux:

~/miniconda3/lib/python3.8/site-packages/jax/api.py in _vjp_pullback_wrapper(cotangent_dtypes, io_tree, fun, py_args)
   1458              "match type of corresponding primal output ({})")
   1459       raise TypeError(msg.format(_dtype(a), dtype))
-&gt; 1460   ans = fun(*args)
   1461   return tree_unflatten(out_tree, ans)
   1462 

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in unbound_vjp(pvals, jaxpr, consts, *cts)
    115     cts = tuple(map(ignore_consts, cts, pvals))
    116     dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
--&gt; 117     arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
    118     return map(instantiate_zeros, arg_cts)
    119 

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
    200         cts_in_avals = [v.aval for v in eqn.outvars]
    201         call_jaxpr, params = core.extract_call_jaxpr(eqn.primitive, eqn.params)
--&gt; 202         cts_out = get_primitive_transpose(eqn.primitive)(
    203             params, call_jaxpr, invals, cts_in, cts_in_avals)
    204       else:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in call_transpose(primitive, params, call_jaxpr, args, ct, _)
    486     new_params = update_params(new_params, map(is_undefined_primal, args),
    487                                [type(x) is not Zero for x in ct])
--&gt; 488   out_flat = primitive.bind(fun, *all_args, **new_params)
    489   return tree_unflatten(out_tree(), out_flat)
    490 primitive_transposes[core.call_p] = partial(call_transpose, call_p)

~/miniconda3/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1132 
   1133   def bind(self, fun, *args, **params):
-&gt; 1134     return call_bind(self, fun, *args, **params)
   1135 
   1136   def process(self, trace, fun, tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1124   else:
   1125     tracers = map(top_trace.full_raise, args)
-&gt; 1126     outs = primitive.process(top_trace, fun, tracers, params)
   1127   return apply_todos(env_trace_todo(), map(full_lower, outs))
   1128 

~/miniconda3/lib/python3.8/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1135 
   1136   def process(self, trace, fun, tracers, params):
-&gt; 1137     return trace.process_call(self, fun, tracers, params)
   1138 
   1139   def post_process(self, trace, out_tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in process_call(self, call_primitive, f, tracers, params)
    273     update_params = call_param_updaters.get(call_primitive)
    274     new_params = update_params(params, nz_tangents) if update_params else params
--&gt; 275     result = call_primitive.bind(f_jvp, *primals, *nonzero_tangents, **new_params)
    276     primal_out, tangent_out = tree_unflatten(out_tree_def(), result)
    277     return [JVPTracer(self, p, t) for p, t in zip(primal_out, tangent_out)]

~/miniconda3/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1132 
   1133   def bind(self, fun, *args, **params):
-&gt; 1134     return call_bind(self, fun, *args, **params)
   1135 
   1136   def process(self, trace, fun, tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1124   else:
   1125     tracers = map(top_trace.full_raise, args)
-&gt; 1126     outs = primitive.process(top_trace, fun, tracers, params)
   1127   return apply_todos(env_trace_todo(), map(full_lower, outs))
   1128 

~/miniconda3/lib/python3.8/site-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1135 
   1136   def process(self, trace, fun, tracers, params):
-&gt; 1137     return trace.process_call(self, fun, tracers, params)
   1138 
   1139   def post_process(self, trace, out_tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in process_call(self, primitive, f, tracers, params)
    179                   else PartialVal.unknown(mapped_aval(pval[0]))
    180                   for pval, is_mapped in zip(in_pvals, params['mapped_invars'])]
--&gt; 181     jaxpr, out_pvals, consts, env_tracers = self.partial_eval(
    182         f, in_pvals, partial(primitive.bind, **params))
    183     if primitive.map_primitive:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in partial_eval(self, f, pvals, app)
    279     f = trace_to_subjaxpr(f, self.master, False)
    280     f, aux = partial_eval_wrapper(f, tuple(in_avals))
--&gt; 281     out_flat, (out_avals, jaxpr, env) = app(f, *in_consts), aux()
    282     out_consts, consts = split_list(out_flat, [len(out_flat)-len(jaxpr.constvars)])
    283     out_pvs = map(PartialVal, zip(out_avals, out_consts))

~/miniconda3/lib/python3.8/site-packages/jax/core.py in bind(self, fun, *args, **params)
   1132 
   1133   def bind(self, fun, *args, **params):
-&gt; 1134     return call_bind(self, fun, *args, **params)
   1135 
   1136   def process(self, trace, fun, tracers, params):

~/miniconda3/lib/python3.8/site-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1121   if top_trace is None:
   1122     with new_sublevel():
-&gt; 1123       outs = primitive.impl(fun, *args, **params)
   1124   else:
   1125     tracers = map(top_trace.full_raise, args)

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
    524 
    525 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
--&gt; 526   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
    527                                *unsafe_map(arg_spec, args))
    528   try:

~/miniconda3/lib/python3.8/site-packages/jax/linear_util.py in memoized_fun(fun, *args)
    222       fun.populate_stores(stores)
    223     else:
--&gt; 224       ans = call(fun, *args)
    225       cache[key] = (ans, fun.stores)
    226     return ans

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    595   else:
    596     pvals: Sequence[pe.PartialVal] = [pe.PartialVal.unknown(aval) for aval in abstract_args]
--&gt; 597     jaxpr, pvals, consts = pe.trace_to_jaxpr(
    598         fun, pvals, instantiate=False, stage_out=True, bottom=True)
    599   map(prefetch, it.chain(consts, jaxpr_literals(jaxpr)))

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out, bottom, trace_type)
    421   with core.new_master(trace_type, bottom=bottom) as master:
    422     fun = trace_to_subjaxpr(fun, master, instantiate)
--&gt; 423     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    424     assert not env
    425     del master

~/miniconda3/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    148     gen = None
    149 
--&gt; 150     ans = self.f(*args, **dict(self.params, **kwargs))
    151     del args
    152     while stack:

~/miniconda3/lib/python3.8/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
    203             params, call_jaxpr, invals, cts_in, cts_in_avals)
    204       else:
--&gt; 205         cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
    206                                                          **eqn.params)
    207     cts_out = [Zero(v.aval) for v in eqn.invars] if cts_out is Zero else cts_out

~/miniconda3/lib/python3.8/site-packages/jax/lax/lax_control_flow.py in _while_transpose_error(*_, **kwargs)
    536 
    537 def _while_transpose_error(*_, **kwargs):
--&gt; 538   raise ValueError("Reverse-mode differentiation does not work for "
    539                    "lax.while_loop or lax.fori_loop. "
    540                    "Try using lax.scan instead.")

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead.
</code></pre></div>
<h3 dir="auto">Description</h3>
<p dir="auto">When we instantiate a BCOO array, there are situations in which we might like its <code class="notranslate">shape</code> to be a <code class="notranslate">DeviceArray</code> (or sequence of <code class="notranslate">DeviceArray</code>s). In particular, we might like the <code class="notranslate">shape</code> to be returned by a compiled function. JAX doesn't complain when we do this — great.</p>
<p dir="auto">However, if we call the <code class="notranslate">BCOO</code> constructor inside a <code class="notranslate">jit</code>-compiled function with a <code class="notranslate">DeviceArray</code>-valued <code class="notranslate">shape</code>, it results in a tracer leak. Here's a minimal example:</p>
<h3 dir="auto">Setup</h3>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
from jax.experimental.sparse import BCOO

mk_bcoo = lambda data, indices, shape: BCOO((data, indices), shape=shape)

data = jnp.array([1, 1])
indices = jnp.array([[0], [1]])
shape = jnp.array([99])"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span>.<span class="pl-s1">sparse</span> <span class="pl-k">import</span> <span class="pl-v">BCOO</span>

<span class="pl-s1">mk_bcoo</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">data</span>, <span class="pl-s1">indices</span>, <span class="pl-s1">shape</span>: <span class="pl-v">BCOO</span>((<span class="pl-s1">data</span>, <span class="pl-s1">indices</span>), <span class="pl-s1">shape</span><span class="pl-c1">=</span><span class="pl-s1">shape</span>)

<span class="pl-s1">data</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>])
<span class="pl-s1">indices</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([[<span class="pl-c1">0</span>], [<span class="pl-c1">1</span>]])
<span class="pl-s1">shape</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">99</span>])</pre></div>
<h3 dir="auto">Results</h3>
<p dir="auto">It works as expected when we don't use <code class="notranslate">jit</code> compilation.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt; mk_bcoo(data, indices, shape)
BCOO(int32[DeviceArray(99, dtype=int32)], nse=2)

&gt;&gt; mk_bcoo(data, indices, shape).todense()
# Works!"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span> <span class="pl-en">mk_bcoo</span>(<span class="pl-s1">data</span>, <span class="pl-s1">indices</span>, <span class="pl-s1">shape</span>)
<span class="pl-v">BCOO</span>(<span class="pl-s1">int32</span>[<span class="pl-v">DeviceArray</span>(<span class="pl-c1">99</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">int32</span>)], <span class="pl-s1">nse</span><span class="pl-c1">=</span><span class="pl-c1">2</span>)

<span class="pl-c1">&gt;&gt;</span> <span class="pl-en">mk_bcoo</span>(<span class="pl-s1">data</span>, <span class="pl-s1">indices</span>, <span class="pl-s1">shape</span>).<span class="pl-en">todense</span>()
<span class="pl-c"># Works!</span></pre></div>
<p dir="auto">But if we <code class="notranslate">jit</code>, the tracer for the shape is leaked.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt; jax.jit(mk_bcoo)(data, indices, shape)
BCOO(int32[Traced&lt;ShapedArray(int32[])&gt;with&lt;DynamicJaxprTrace(level=0/1)&gt;], nse=2)

&gt;&gt; result = jax.jit(mk_bcoo)(data, indices, shape); result.shape
(Traced&lt;ShapedArray(int32[])&gt;with&lt;DynamicJaxprTrace(level=0/1)&gt;,)

&gt;&gt; jax.jit(mk_bcoo)(data, indices, shape).todense()
# Raises a TracerArrayConversionError"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">mk_bcoo</span>)(<span class="pl-s1">data</span>, <span class="pl-s1">indices</span>, <span class="pl-s1">shape</span>)
<span class="pl-v">BCOO</span>(<span class="pl-s1">int32</span>[<span class="pl-v">Traced</span><span class="pl-c1">&lt;</span><span class="pl-v">ShapedArray</span>(<span class="pl-s1">int32</span>[<span class="pl-s1"></span>])<span class="pl-c1">&gt;</span><span class="pl-s1">with</span><span class="pl-c1">&lt;</span><span class="pl-v">DynamicJaxprTrace</span>(<span class="pl-s1">level</span><span class="pl-c1">=</span><span class="pl-c1">0</span><span class="pl-c1">/</span><span class="pl-c1">1</span>)<span class="pl-c1">&gt;</span><span class="pl-s1"></span>], <span class="pl-s1">nse</span><span class="pl-c1">=</span><span class="pl-c1">2</span>)

<span class="pl-c1">&gt;&gt;</span> <span class="pl-s1">result</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">mk_bcoo</span>)(<span class="pl-s1">data</span>, <span class="pl-s1">indices</span>, <span class="pl-s1">shape</span>); <span class="pl-s1">result</span>.<span class="pl-s1">shape</span>
(<span class="pl-v">Traced</span><span class="pl-c1">&lt;</span><span class="pl-v">ShapedArray</span>(<span class="pl-s1">int32</span>[<span class="pl-s1"></span>])<span class="pl-c1">&gt;</span><span class="pl-s1">with</span><span class="pl-c1">&lt;</span><span class="pl-v">DynamicJaxprTrace</span>(<span class="pl-s1">level</span><span class="pl-c1">=</span><span class="pl-c1">0</span><span class="pl-c1">/</span><span class="pl-c1">1</span>)<span class="pl-c1">&gt;</span><span class="pl-s1"></span>,)

<span class="pl-c1">&gt;&gt;</span> <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">mk_bcoo</span>)(<span class="pl-s1">data</span>, <span class="pl-s1">indices</span>, <span class="pl-s1">shape</span>).<span class="pl-en">todense</span>()
<span class="pl-c"># Raises a TracerArrayConversionError</span></pre></div>
<p dir="auto">Apologies if this is a duplicate.</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">JAX v0.3.16</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU</p>
<h3 dir="auto">Additional System Info</h3>
<p dir="auto"><em>No response</em></p>