Issue 1:
Would it be possible to implement this interesting paper idea with jax / flax?

Intelligent Matrix Exponentiation

paper: https://arxiv.org/pdf/2008.03936.pdf

code: https://github.com/google-research/google-research/tree/master/m_layer

wiki: https://en.wikipedia.org/wiki/Matrix_exponential

expm docs: https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.expm.html#jax.scipy.linalg.expm

expm frechet looks like the right thing to make the gradient:

https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.linalg.expm_frechet.html#jax.scipy.linalg.expm_frechet

I tried a few options, here's a simple one:

import jax
from flax import nn
jnp = jax.numpy
vec_expm = jnp.vectorize(jax.scipy.linalg.expm, signature='(k)->(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

however this crashes because of this:

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead.

Would it work if we wire up expm_frechet primitive?

a simpler reproduction without other nn stuff:

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)

full trace:

---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-14-f2d123ee1505> in <module>
----> 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:
--> 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)
-> 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:
--> 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)
---> 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)
--> 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
--> 150 ans = self.f(*args, **dict(self.params, **kwargs))
151 del args
152 while stack:

<ipython-input-10-995182967638> in wrapped(*args, **kwargs)
1 def value_and_jacobian(fun):
2 def wrapped(*args, **kwargs):
----> 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)
--> 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))
-> 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]
--> 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)
--> 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])
--> 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):
-> 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)
-> 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):
-> 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
--> 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):
-> 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)
-> 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):
-> 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'])]
--> 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))
--> 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):
-> 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():
-> 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):
--> 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:
--> 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]
--> 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)
--> 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
--> 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:
--> 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):
--> 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.

Issue 2:
Description

When we instantiate a BCOO array, there are situations in which we might like its shape to be a DeviceArray (or sequence of DeviceArrays). In particular, we might like the shape to be returned by a compiled function. JAX doesn't complain when we do this — great.

However, if we call the BCOO constructor inside a jit-compiled function with a DeviceArray-valued shape, it results in a tracer leak. Here's a minimal example:

Setup

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])

Results

It works as expected when we don't use jit compilation.

>> mk_bcoo(data, indices, shape)
BCOO(int32[DeviceArray(99, dtype=int32)], nse=2)

>> mk_bcoo(data, indices, shape).todense()
# Works!

But if we jit, the tracer for the shape is leaked.

>> jax.jit(mk_bcoo)(data, indices, shape)
BCOO(int32[Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>], nse=2)

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

>> jax.jit(mk_bcoo)(data, indices, shape).todense()
# Raises a TracerArrayConversionError

Apologies if this is a duplicate.

What jax/jaxlib version are you using?

JAX v0.3.16

Which accelerator(s) are you using?

CPU

Additional System Info

No response