Issue 1:
Description

Consider these 2 equivalent functions. Only the one that doesn't use scan is differentiable via jax.linear_transpose:

def foo(p):
x = jnp.zeros_like(p)
for _ in range(10):
x += p
return x

def bar(p):
return jax.lax.scan(lambda x, _: (x + p, None), jnp.zeros_like(p), None, length=10)[0]

# works
jax.linear_transpose(foo, jnp.zeros(10))(jnp.ones(10))

# nope
jax.linear_transpose(bar, jnp.zeros(10))(jnp.ones(10))

Traceback:

---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation Traceback (most recent call last)
File /opt/homebrew/Cellar/python@3.10/3.10.5/Frameworks/Python.framework/Versions/3.10/lib/python3.10/runpy.py:191, in _run_module_as_main(***failed resolving arguments***)
190 except _Error as exc:
--> 191 msg = "%s: %s" % (sys.executable, exc)
192 sys.exit(msg)

File /opt/homebrew/Cellar/python@3.10/3.10.5/Frameworks/Python.framework/Versions/3.10/lib/python3.10/runpy.py:75, in _run_code(***failed resolving arguments***)
74 loader = mod_spec.loader
---> 75 fname = mod_spec.origin
76 cached = mod_spec.cached

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel_launcher.py:12, in <module>
9 if __name__ == "__main__":
10 # Remove the CWD from sys.path while we load stuff.
11 # This is added back by InteractiveShellApp.init_path()
---> 12 if sys.path[0] == "":
13 del sys.path[0]

File ~/.virtualenvs/science/lib/python3.10/site-packages/traitlets/config/application.py:963, in Application.launch_instance(***failed resolving arguments***)
959 """Launch a global instance of this Application
960
961 If a global instance already exists, this reinitializes and starts it
962 """
--> 963 app = cls.instance(**kwargs)
964 app.initialize(argv)

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel/kernelapp.py:702, in IPKernelApp.start(***failed resolving arguments***)
701 if self.trio_loop:
--> 702 from ipykernel.trio_runner import TrioRunner
704 tr = TrioRunner()

File ~/.virtualenvs/science/lib/python3.10/site-packages/tornado/platform/asyncio.py:195, in BaseAsyncIOLoop.start(***failed resolving arguments***)
194 except (RuntimeError, AssertionError):
--> 195 old_loop = None # type: ignore
196 try:

File /opt/homebrew/Cellar/python@3.10/3.10.5/Frameworks/Python.framework/Versions/3.10/lib/python3.10/asyncio/base_events.py:594, in BaseEventLoop.run_forever(***failed resolving arguments***)
592 self._thread_id = threading.get_ident()
--> 594 old_agen_hooks = sys.get_asyncgen_hooks()
595 sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
596 finalizer=self._asyncgen_finalizer_hook)

File /opt/homebrew/Cellar/python@3.10/3.10.5/Frameworks/Python.framework/Versions/3.10/lib/python3.10/asyncio/base_events.py:1860, in BaseEventLoop._run_once(***failed resolving arguments***)
1858 timeout = min(max(0, when - self.time()), MAXIMUM_SELECT_TIMEOUT)
-> 1860 event_list = self._selector.select(timeout)
1861 self._process_events(event_list)

File /opt/homebrew/Cellar/python@3.10/3.10.5/Frameworks/Python.framework/Versions/3.10/lib/python3.10/asyncio/events.py:80, in Handle._run(***failed resolving arguments***)
79 try:
---> 80 self._context.run(self._callback, *self._args)
81 except (SystemExit, KeyboardInterrupt):

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel/kernelbase.py:504, in Kernel.dispatch_queue(***failed resolving arguments***)
503 try:
--> 504 await self.process_one()
505 except Exception:

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel/kernelbase.py:490, in Kernel.process_one(***failed resolving arguments***)
489 try:
--> 490 t, dispatch, args = self.msg_queue.get_nowait()
491 except (asyncio.QueueEmpty, QueueEmpty):

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel/kernelbase.py:377, in Kernel.dispatch_shell(***failed resolving arguments***)
376 self.shell_stream.flush(zmq.POLLOUT)
--> 377 return
379 # Print some info about this message and leave a '--->' marker, so it's
380 # easier to trace visually the message chain when debugging. Each
381 # handler prints its message at the end.

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel/kernelbase.py:695, in Kernel.execute_request(***failed resolving arguments***)
693 stop_on_error = content.get("stop_on_error", True)
--> 695 metadata = self.init_metadata(parent)
697 # Re-broadcast our input for the benefit of listening clients, and
698 # start computing output

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel/ipkernel.py:352, in IPythonKernel.do_execute(***failed resolving arguments***)
350 if with_cell_id:
351 coro = run_cell(
--> 352 code,
353 store_history=store_history,
354 silent=silent,
355 transformed_cell=transformed_cell,
356 preprocessing_exc_tuple=preprocessing_exc_tuple,
357 cell_id=cell_id,
358 )
359 else:

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel/zmqshell.py:528, in ZMQInteractiveShell.run_cell(***failed resolving arguments***)
527 self._last_traceback = None
--> 528 return super().run_cell(*args, **kwargs)

File ~/.virtualenvs/science/lib/python3.10/site-packages/IPython/core/interactiveshell.py:2881, in InteractiveShell.run_cell(***failed resolving arguments***)
2879 try:
2880 result = self._run_cell(
-> 2881 raw_cell, store_history, silent, shell_futures, cell_id
2882 )
2883 finally:

File ~/.virtualenvs/science/lib/python3.10/site-packages/IPython/core/interactiveshell.py:2910, in InteractiveShell._run_cell(***failed resolving arguments***)
2908 assert transformed_cell is not None
2909 coro = self.run_cell_async(
-> 2910 raw_cell,
2911 store_history=store_history,
2912 silent=silent,
2913 shell_futures=shell_futures,
2914 transformed_cell=transformed_cell,
2915 preprocessing_exc_tuple=preprocessing_exc_tuple,
2916 cell_id=cell_id,
2917 )
2919 # run_cell_async is async, but may not actually need an eventloop.
2920 # when this is the case, we want to run it using the pseudo_sync_runner
2921 # so that code can invoke eventloops (for example via the %run , and
2922 # `%paste` magic.

File ~/.virtualenvs/science/lib/python3.10/site-packages/IPython/core/async_helpers.py:129, in _pseudo_sync_runner(***failed resolving arguments***)
128 try:
--> 129 coro.send(None)
130 except StopIteration as exc:

File ~/.virtualenvs/science/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3108, in InteractiveShell.run_cell_async(***failed resolving arguments***)
3107 code_ast = compiler.ast_parse(cell, filename=cell_name)
-> 3108 except self.custom_exceptions as e:
3109 etype, value, tb = sys.exc_info()

File ~/.virtualenvs/science/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3305, in InteractiveShell.run_ast_nodes(***failed resolving arguments***)
3304 to_run_exec, to_run_interactive = nodelist[:-1], nodelist[-1:]
-> 3305 elif interactivity == 'all':
3306 to_run_exec, to_run_interactive = [], nodelist

File ~/.virtualenvs/science/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3394, in InteractiveShell.run_code(***failed resolving arguments***)
3393 try:
-> 3394 if async_:
3395 await eval(code_obj, self.user_global_ns, self.user_ns)

Input In [207], in <cell line: 10>()
8 return x
---> 10 jax.linear_transpose(foo, jnp.zeros(10))(jnp.ones(10))

Input In [207], in foo(***failed resolving arguments***)
1 def foo(p):
----> 2 return jax.lax.scan(lambda x, _: (x + p, None), jnp.zeros_like(p), None, length=10)[0]

JaxStackTraceBeforeTransformation: AssertionError

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

AssertionError Traceback (most recent call last)
Input In [207], in <cell line: 10>()
7 x += p
8 return x
---> 10 jax.linear_transpose(foo, jnp.zeros(10))(jnp.ones(10))

[... skipping hidden 3 frame]

File ~/.virtualenvs/science/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py:663, in _scan_transpose(reduce_axes, cts, reverse, length, num_consts, num_carry, jaxpr, linear, unroll, *args)
661 ires, _ = split_list(consts, [num_ires])
662 _, eres = split_list(xs, [sum(xs_lin)])
--> 663 assert not any(ad.is_undefined_primal(r) for r in ires)
664 assert not any(ad.is_undefined_primal(r) for r in eres)
666 carry_avals, y_avals = split_list(jaxpr.out_avals, [num_carry])

AssertionError:

What jax/jaxlib version are you using?

v0.3.15

Which accelerator(s) are you using?

CPU

Additional System Info

OSX ARM

Issue 2:
first some motivation:

I have a hermitian linear operator which is using a jvp followed by a vjp of complex (shallow) neural networks which ~ looks like this:

def matvec(apply_fun, params, samples, v):
_, w = jax.jvp(lambda p: apply_fun(p, samples), (params,), (v,))
w = w.conj()
res, _ = jax.vjp(apply_fun, params, samples)[1](w)
return jax.tree_map(jax.lax.conj, res)

It works fine, however storing the forward pass of the vjp uses too much memory when using many parameters and samples (think of both like ~1M). However one can to split the vjp into smaller vjp's looping over batches of samples and summing the result (which works surprisingly well). See the snipped below which should make this clear.

I am using this linear operator with the cg and gmres iterative solvers, where its required to be linear_transpose'able.

I have implemented the loop with lax.scan to try to make it transposable, however the _scan_transpose rule is giving an error which I dont't really know how to solve.

Here is a minimal snipped of code (where for simplicity I have replaced the jvp/vjp in the linear operator with products with the matrix J), which reproduces the error I am getting:

import jax
import jax.numpy as jnp
from functools import partial

n = 100
batchsize = 5
d = 10

# generate a random jacobian and a vector
k0, k1, k2, k3 = jax.random.split(jax.random.PRNGKey(0), 4)
J = jax.random.uniform(k0, shape=(n,d)) + 1j * jax.random.uniform(k1, shape=(n,d))
v = jax.random.uniform(k2, shape=(d,)) + 1j * jax.random.normal(k3, shape=(d,))

# the linar operator
def mv(J, v):
return J.T.conj()@(J@v)

# works fine
_ = jax.scipy.sparse.linalg.cg(partial(mv, J), v)
_ = jax.scipy.sparse.linalg.gmres(partial(mv, J), v, solve_method='incremental')

# mv is hermitian, so its transpose is trivial
def mv_transposed(J, v):
return mv(J, v.conj()).conj()

# now the operator where the vjp is looping over small batches
def mv_batched(J, v):
w = J@v

J_batched = J.reshape((-1, batchsize)+J.shape[1:])
w_batched = w.reshape((-1, batchsize)+w.shape[1:])
Jw = J_batched, w_batched

def _mv(Jw):
J, w = Jw
return (w.conj() @ J).conj()

def f(carry, x):
return carry + _mv(x), None

res, _ = jax.lax.scan(f, jnp.zeros_like(v), Jw, unroll=1)

return res

assert jnp.allclose(mv(J,v), mv_batched(J,v))
assert jnp.allclose(mv_transposed(J,v), jax.linear_transpose(partial(mv, J), v)(v)[0])

# linear_transposing mv_batched gives the error:
print( jax.linear_transpose(partial(mv_batched, J), v)(v)[0] )

# goal: use mv_batched with the iterative solvers
# same error:
# print( jax.scipy.sparse.linalg.cg(partial(mv_batched, J), v))
# print( jax.scipy.sparse.linalg.gmres(partial(mv_batched, J), v, solve_method='incremental'))

And the output:

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
---------------------------------------------------------------------------
AssertionError Traceback (most recent call last)
<ipython-input-1-dbeaca6d614d> in <module>
45 # jax.scipy.sparse.linalg.gmres(partial(mv_batched, J), v, solve_method='incremental')
46
---> 47 print( jax.linear_transpose(partial(mv_batched, J), v)(v)[0] )

/usr/local/lib/python3.9/site-packages/jax/api.py in transposed_fun(out_cotangent)
1957 in_cotangents = map(
1958 ad.instantiate_zeros,
-> 1959 ad.backward_pass(jaxpr, consts, dummies, out_cotangents))
1960 return tree_unflatten(in_tree, in_cotangents)
1961

/usr/local/lib/python3.9/site-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
216 params, call_jaxpr, invals, cts_in, cts_in_avals)
217 else:
--> 218 cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
219 **eqn.params)
220 cts_out = [Zero(v.aval) for v in eqn.invars] if cts_out is Zero else cts_out

/usr/local/lib/python3.9/site-packages/jax/_src/lax/control_flow.py in _scan_transpose(cts, reverse, length, num_consts, num_carry, jaxpr, linear, unroll, *args)
1680 _, eres = split_list(xs, [sum(xs_lin)])
1681 assert not any(ad.is_undefined_primal(r) for r in ires)
-> 1682 assert not any(ad.is_undefined_primal(r) for r in eres)
1683
1684 carry_avals, y_avals = split_list(jaxpr.out_avals, [num_carry])

AssertionError:

[I thought of just using linear_call, however it's still missing a XLA translation rule so it currently can't be jitted]

Please see this more like a feature request for a more universal transpose rule (unless there's a trivial trick to make it transposable which I'm missing; otherwise I would also be happy to have a go at trying to relax the requirement of the linear operator being transposable in the solvers if you point me in the right direction since there are several possible ways to do this).