<h3 dir="auto">Description</h3>
<p dir="auto">Consider these 2 equivalent functions. Only the one that doesn't use <code class="notranslate">scan</code> is differentiable via <code class="notranslate">jax.linear_transpose</code>:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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))"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">p</span>):
    <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">zeros_like</span>(<span class="pl-s1">p</span>)
    <span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">10</span>):
        <span class="pl-s1">x</span> <span class="pl-c1">+=</span> <span class="pl-s1">p</span>
    <span class="pl-k">return</span> <span class="pl-s1">x</span>

<span class="pl-k">def</span> <span class="pl-en">bar</span>(<span class="pl-s1">p</span>):
    <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">scan</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>, <span class="pl-s1">_</span>: (<span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-s1">p</span>, <span class="pl-c1">None</span>), <span class="pl-s1">jnp</span>.<span class="pl-en">zeros_like</span>(<span class="pl-s1">p</span>), <span class="pl-c1">None</span>, <span class="pl-s1">length</span><span class="pl-c1">=</span><span class="pl-c1">10</span>)[<span class="pl-c1">0</span>]

<span class="pl-c"># works</span>
<span class="pl-s1">jax</span>.<span class="pl-en">linear_transpose</span>(<span class="pl-s1">foo</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">10</span>))(<span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">10</span>))

<span class="pl-c"># nope</span>
<span class="pl-s1">jax</span>.<span class="pl-en">linear_transpose</span>(<span class="pl-s1">bar</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">10</span>))(<span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">10</span>))</pre></div>
<p dir="auto">Traceback:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
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:
--&gt; 191     msg = &quot;%s: %s&quot; % (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
---&gt; 75 fname = mod_spec.origin
     76 cached = mod_spec.cached

File ~/.virtualenvs/science/lib/python3.10/site-packages/ipykernel_launcher.py:12, in &lt;module&gt;
      9 if __name__ == &quot;__main__&quot;:
     10     # Remove the CWD from sys.path while we load stuff.
     11     # This is added back by InteractiveShellApp.init_path()
---&gt; 12     if sys.path[0] == &quot;&quot;:
     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 &quot;&quot;&quot;Launch a global instance of this Application
    960 
    961 If a global instance already exists, this reinitializes and starts it
    962 &quot;&quot;&quot;
--&gt; 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:
--&gt; 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):
--&gt; 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()
--&gt; 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)
-&gt; 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:
---&gt; 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:
--&gt; 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:
--&gt; 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)
--&gt; 377     return
    379 # Print some info about this message and leave a '---&gt;' 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(&quot;stop_on_error&quot;, True)
--&gt; 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(
--&gt; 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
--&gt; 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(
-&gt; 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(
-&gt; 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:
--&gt; 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)
-&gt; 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:]
-&gt; 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:
-&gt; 3394     if async_:
   3395         await eval(code_obj, self.user_global_ns, self.user_ns)

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

Input In [207], in foo(***failed resolving arguments***)
      1 def foo(p):
----&gt; 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 &lt;cell line: 10&gt;()
      7         x += p
      8     return x
---&gt; 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)])
--&gt; 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: "><pre class="notranslate"><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span>
<span class="pl-v">JaxStackTraceBeforeTransformation</span>         <span class="pl-v">Traceback</span> (<span class="pl-s1">most</span> <span class="pl-s1">recent</span> <span class="pl-s1">call</span> <span class="pl-s1">last</span>)
<span class="pl-v">File</span> <span class="pl-c1">/</span><span class="pl-s1">opt</span><span class="pl-c1">/</span><span class="pl-s1">homebrew</span><span class="pl-c1">/</span><span class="pl-v">Cellar</span><span class="pl-c1">/</span><span class="pl-s1">python</span>@<span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span>.<span class="pl-c1">5</span><span class="pl-c1">/</span><span class="pl-v">Frameworks</span><span class="pl-c1">/</span><span class="pl-v">Python</span>.<span class="pl-s1">framework</span><span class="pl-c1">/</span><span class="pl-v">Versions</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">runpy</span>.<span class="pl-s1">py</span>:<span class="pl-c1">191</span>, <span class="pl-s1">in</span> <span class="pl-en">_run_module_as_main</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">190</span> <span class="pl-s1">except</span> <span class="pl-s1">_Error</span> <span class="pl-k">as</span> <span class="pl-s1">exc</span>:
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">191</span>     <span class="pl-s1">msg</span> <span class="pl-c1">=</span> <span class="pl-s">"%s: %s"</span> <span class="pl-c1">%</span> (<span class="pl-s1">sys</span>.<span class="pl-s1">executable</span>, <span class="pl-s1">exc</span>)
    <span class="pl-c1">192</span>     <span class="pl-s1">sys</span>.<span class="pl-en">exit</span>(<span class="pl-s1">msg</span>)

<span class="pl-v">File</span> <span class="pl-c1">/</span><span class="pl-s1">opt</span><span class="pl-c1">/</span><span class="pl-s1">homebrew</span><span class="pl-c1">/</span><span class="pl-v">Cellar</span><span class="pl-c1">/</span><span class="pl-s1">python</span>@<span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span>.<span class="pl-c1">5</span><span class="pl-c1">/</span><span class="pl-v">Frameworks</span><span class="pl-c1">/</span><span class="pl-v">Python</span>.<span class="pl-s1">framework</span><span class="pl-c1">/</span><span class="pl-v">Versions</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">runpy</span>.<span class="pl-s1">py</span>:<span class="pl-c1">75</span>, <span class="pl-s1">in</span> <span class="pl-en">_run_code</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
     <span class="pl-c1">74</span> <span class="pl-s1">loader</span> <span class="pl-c1">=</span> <span class="pl-s1">mod_spec</span>.<span class="pl-s1">loader</span>
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">75</span> <span class="pl-s1">fname</span> <span class="pl-c1">=</span> <span class="pl-s1">mod_spec</span>.<span class="pl-s1">origin</span>
     <span class="pl-c1">76</span> <span class="pl-s1">cached</span> <span class="pl-c1">=</span> <span class="pl-s1">mod_spec</span>.<span class="pl-s1">cached</span>

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

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

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">ipykernel</span><span class="pl-c1">/</span><span class="pl-s1">kernelapp</span>.<span class="pl-s1">py</span>:<span class="pl-c1">702</span>, <span class="pl-s1">in</span> <span class="pl-v">IPKernelApp</span>.<span class="pl-s1">start</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">701</span> <span class="pl-k">if</span> <span class="pl-s1">self</span>.<span class="pl-s1">trio_loop</span>:
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">702</span>     <span class="pl-s1">from</span> <span class="pl-s1">ipykernel</span>.<span class="pl-s1">trio_runner</span> <span class="pl-s1">import</span> <span class="pl-v">TrioRunner</span>
    <span class="pl-c1">704</span>     <span class="pl-s1">tr</span> <span class="pl-c1">=</span> <span class="pl-v">TrioRunner</span>()

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">tornado</span><span class="pl-c1">/</span><span class="pl-s1">platform</span><span class="pl-c1">/</span><span class="pl-s1">asyncio</span>.<span class="pl-s1">py</span>:<span class="pl-c1">195</span>, <span class="pl-c1">in</span> <span class="pl-v">BaseAsyncIOLoop</span>.<span class="pl-s1">start</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">194</span> <span class="pl-en">except</span> (<span class="pl-v">RuntimeError</span>, <span class="pl-v">AssertionError</span>):
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">195</span>     <span class="pl-s1">old_loop</span> <span class="pl-c1">=</span> <span class="pl-c1">None</span>  <span class="pl-c"># type: ignore</span>
    <span class="pl-c1">196</span> <span class="pl-s1">try</span>:

<span class="pl-v">File</span> <span class="pl-c1">/</span><span class="pl-s1">opt</span><span class="pl-c1">/</span><span class="pl-s1">homebrew</span><span class="pl-c1">/</span><span class="pl-v">Cellar</span><span class="pl-c1">/</span><span class="pl-s1">python</span>@<span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span>.<span class="pl-c1">5</span><span class="pl-c1">/</span><span class="pl-v">Frameworks</span><span class="pl-c1">/</span><span class="pl-v">Python</span>.<span class="pl-s1">framework</span><span class="pl-c1">/</span><span class="pl-v">Versions</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">asyncio</span><span class="pl-c1">/</span><span class="pl-s1">base_events</span>.<span class="pl-s1">py</span>:<span class="pl-c1">594</span>, <span class="pl-s1">in</span> <span class="pl-v">BaseEventLoop</span>.<span class="pl-s1">run_forever</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">592</span> <span class="pl-s1">self</span>.<span class="pl-s1">_thread_id</span> <span class="pl-c1">=</span> <span class="pl-s1">threading</span>.<span class="pl-en">get_ident</span>()
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">594</span> <span class="pl-s1">old_agen_hooks</span> <span class="pl-c1">=</span> <span class="pl-s1">sys</span>.<span class="pl-en">get_asyncgen_hooks</span>()
    <span class="pl-c1">595</span> <span class="pl-s1">sys</span>.<span class="pl-en">set_asyncgen_hooks</span>(<span class="pl-s1">firstiter</span><span class="pl-c1">=</span><span class="pl-s1">self</span>.<span class="pl-s1">_asyncgen_firstiter_hook</span>,
    <span class="pl-c1">596</span>                        <span class="pl-s1">finalizer</span><span class="pl-c1">=</span><span class="pl-s1">self</span>.<span class="pl-s1">_asyncgen_finalizer_hook</span>)

<span class="pl-v">File</span> <span class="pl-c1">/</span><span class="pl-s1">opt</span><span class="pl-c1">/</span><span class="pl-s1">homebrew</span><span class="pl-c1">/</span><span class="pl-v">Cellar</span><span class="pl-c1">/</span><span class="pl-s1">python</span>@<span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span>.<span class="pl-c1">5</span><span class="pl-c1">/</span><span class="pl-v">Frameworks</span><span class="pl-c1">/</span><span class="pl-v">Python</span>.<span class="pl-s1">framework</span><span class="pl-c1">/</span><span class="pl-v">Versions</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">asyncio</span><span class="pl-c1">/</span><span class="pl-s1">base_events</span>.<span class="pl-s1">py</span>:<span class="pl-c1">1860</span>, <span class="pl-s1">in</span> <span class="pl-v">BaseEventLoop</span>.<span class="pl-s1">_run_once</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
   <span class="pl-c1">1858</span>     <span class="pl-s1">timeout</span> <span class="pl-c1">=</span> <span class="pl-en">min</span>(<span class="pl-en">max</span>(<span class="pl-c1">0</span>, <span class="pl-s1">when</span> <span class="pl-c1">-</span> <span class="pl-s1">self</span>.<span class="pl-en">time</span>()), <span class="pl-v">MAXIMUM_SELECT_TIMEOUT</span>)
<span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">1860</span> <span class="pl-s1">event_list</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">_selector</span>.<span class="pl-en">select</span>(<span class="pl-s1">timeout</span>)
   <span class="pl-c1">1861</span> <span class="pl-s1">self</span>.<span class="pl-en">_process_events</span>(<span class="pl-s1">event_list</span>)

<span class="pl-v">File</span> <span class="pl-c1">/</span><span class="pl-s1">opt</span><span class="pl-c1">/</span><span class="pl-s1">homebrew</span><span class="pl-c1">/</span><span class="pl-v">Cellar</span><span class="pl-c1">/</span><span class="pl-s1">python</span>@<span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span>.<span class="pl-c1">5</span><span class="pl-c1">/</span><span class="pl-v">Frameworks</span><span class="pl-c1">/</span><span class="pl-v">Python</span>.<span class="pl-s1">framework</span><span class="pl-c1">/</span><span class="pl-v">Versions</span><span class="pl-c1">/</span><span class="pl-c1">3.10</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">asyncio</span><span class="pl-c1">/</span><span class="pl-s1">events</span>.<span class="pl-s1">py</span>:<span class="pl-c1">80</span>, <span class="pl-s1">in</span> <span class="pl-v">Handle</span>.<span class="pl-s1">_run</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
     <span class="pl-c1">79</span> <span class="pl-s1">try</span>:
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">80</span>     <span class="pl-s1">self</span>.<span class="pl-s1">_context</span>.<span class="pl-en">run</span>(<span class="pl-s1">self</span>.<span class="pl-s1">_callback</span>, <span class="pl-c1">*</span><span class="pl-s1">self</span>.<span class="pl-s1">_args</span>)
     <span class="pl-c1">81</span> <span class="pl-en">except</span> (<span class="pl-v">SystemExit</span>, <span class="pl-v">KeyboardInterrupt</span>):

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">ipykernel</span><span class="pl-c1">/</span><span class="pl-s1">kernelbase</span>.<span class="pl-s1">py</span>:<span class="pl-c1">504</span>, <span class="pl-s1">in</span> <span class="pl-v">Kernel</span>.<span class="pl-s1">dispatch_queue</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">503</span> <span class="pl-s1">try</span>:
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">504</span>     <span class="pl-s1">await</span> <span class="pl-s1">self</span>.<span class="pl-en">process_one</span>()
    <span class="pl-c1">505</span> <span class="pl-s1">except</span> <span class="pl-v">Exception</span>:

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">ipykernel</span><span class="pl-c1">/</span><span class="pl-s1">kernelbase</span>.<span class="pl-s1">py</span>:<span class="pl-c1">490</span>, <span class="pl-s1">in</span> <span class="pl-v">Kernel</span>.<span class="pl-s1">process_one</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">489</span> <span class="pl-s1">try</span>:
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">490</span>     <span class="pl-s1">t</span>, <span class="pl-s1">dispatch</span>, <span class="pl-s1">args</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">msg_queue</span>.<span class="pl-en">get_nowait</span>()
    <span class="pl-c1">491</span> <span class="pl-en">except</span> (<span class="pl-s1">asyncio</span>.<span class="pl-v">QueueEmpty</span>, <span class="pl-v">QueueEmpty</span>):

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

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

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">ipykernel</span><span class="pl-c1">/</span><span class="pl-s1">ipkernel</span>.<span class="pl-s1">py</span>:<span class="pl-c1">352</span>, <span class="pl-s1">in</span> <span class="pl-v">IPythonKernel</span>.<span class="pl-s1">do_execute</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">350</span> <span class="pl-k">if</span> <span class="pl-s1">with_cell_id</span>:
    <span class="pl-c1">351</span>     <span class="pl-s1">coro</span> <span class="pl-c1">=</span> <span class="pl-en">run_cell</span>(
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">352</span>         <span class="pl-s1">code</span>,
    <span class="pl-c1">353</span>         <span class="pl-s1">store_history</span><span class="pl-c1">=</span><span class="pl-s1">store_history</span>,
    <span class="pl-c1">354</span>         <span class="pl-s1">silent</span><span class="pl-c1">=</span><span class="pl-s1">silent</span>,
    <span class="pl-c1">355</span>         <span class="pl-s1">transformed_cell</span><span class="pl-c1">=</span><span class="pl-s1">transformed_cell</span>,
    <span class="pl-c1">356</span>         <span class="pl-s1">preprocessing_exc_tuple</span><span class="pl-c1">=</span><span class="pl-s1">preprocessing_exc_tuple</span>,
    <span class="pl-c1">357</span>         <span class="pl-s1">cell_id</span><span class="pl-c1">=</span><span class="pl-s1">cell_id</span>,
    <span class="pl-c1">358</span>     )
    <span class="pl-c1">359</span> <span class="pl-k">else</span>:

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">ipykernel</span><span class="pl-c1">/</span><span class="pl-s1">zmqshell</span>.<span class="pl-s1">py</span>:<span class="pl-c1">528</span>, <span class="pl-s1">in</span> <span class="pl-v">ZMQInteractiveShell</span>.<span class="pl-s1">run_cell</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">527</span> <span class="pl-s1">self</span>.<span class="pl-s1">_last_traceback</span> <span class="pl-c1">=</span> <span class="pl-c1">None</span>
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">528</span> <span class="pl-s1">return</span> <span class="pl-en">super</span>().<span class="pl-en">run_cell</span>(<span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>)

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-v">IPython</span><span class="pl-c1">/</span><span class="pl-s1">core</span><span class="pl-c1">/</span><span class="pl-s1">interactiveshell</span>.<span class="pl-s1">py</span>:<span class="pl-c1">2881</span>, <span class="pl-s1">in</span> <span class="pl-v">InteractiveShell</span>.<span class="pl-s1">run_cell</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
   <span class="pl-c1">2879</span> <span class="pl-s1">try</span>:
   <span class="pl-c1">2880</span>     <span class="pl-s1">result</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">_run_cell</span>(
<span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">2881</span>         <span class="pl-s1">raw_cell</span>, <span class="pl-s1">store_history</span>, <span class="pl-s1">silent</span>, <span class="pl-s1">shell_futures</span>, <span class="pl-s1">cell_id</span>
   <span class="pl-c1">2882</span>     )
   <span class="pl-c1">2883</span> <span class="pl-k">finally</span>:

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

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-v">IPython</span><span class="pl-c1">/</span><span class="pl-s1">core</span><span class="pl-c1">/</span><span class="pl-s1">async_helpers</span>.<span class="pl-s1">py</span>:<span class="pl-c1">129</span>, <span class="pl-s1">in</span> <span class="pl-s1">_pseudo_sync_runner</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
    <span class="pl-c1">128</span> <span class="pl-s1">try</span>:
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">129</span>     <span class="pl-s1">coro</span>.<span class="pl-en">send</span>(<span class="pl-c1">None</span>)
    <span class="pl-c1">130</span> <span class="pl-s1">except</span> <span class="pl-v">StopIteration</span> <span class="pl-k">as</span> <span class="pl-s1">exc</span>:

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-v">IPython</span><span class="pl-c1">/</span><span class="pl-s1">core</span><span class="pl-c1">/</span><span class="pl-s1">interactiveshell</span>.<span class="pl-s1">py</span>:<span class="pl-c1">3108</span>, <span class="pl-s1">in</span> <span class="pl-v">InteractiveShell</span>.<span class="pl-s1">run_cell_async</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
   <span class="pl-c1">3107</span>     <span class="pl-s1">code_ast</span> <span class="pl-c1">=</span> <span class="pl-s1">compiler</span>.<span class="pl-en">ast_parse</span>(<span class="pl-s1">cell</span>, <span class="pl-s1">filename</span><span class="pl-c1">=</span><span class="pl-s1">cell_name</span>)
<span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">3108</span> <span class="pl-s1">except</span> <span class="pl-s1">self</span>.<span class="pl-s1">custom_exceptions</span> <span class="pl-k">as</span> <span class="pl-s1">e</span>:
   <span class="pl-c1">3109</span>     <span class="pl-s1">etype</span>, <span class="pl-s1">value</span>, <span class="pl-s1">tb</span> <span class="pl-c1">=</span> <span class="pl-s1">sys</span>.<span class="pl-en">exc_info</span>()

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-v">IPython</span><span class="pl-c1">/</span><span class="pl-s1">core</span><span class="pl-c1">/</span><span class="pl-s1">interactiveshell</span>.<span class="pl-s1">py</span>:<span class="pl-c1">3305</span>, <span class="pl-s1">in</span> <span class="pl-v">InteractiveShell</span>.<span class="pl-s1">run_ast_nodes</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
   <span class="pl-c1">3304</span>     <span class="pl-s1">to_run_exec</span>, <span class="pl-s1">to_run_interactive</span> <span class="pl-c1">=</span> <span class="pl-s1">nodelist</span>[:<span class="pl-c1">-</span><span class="pl-c1">1</span>], <span class="pl-s1">nodelist</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>:]
<span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">3305</span> <span class="pl-s1">elif</span> <span class="pl-s1">interactivity</span> <span class="pl-c1">==</span> <span class="pl-s">'all'</span>:
   <span class="pl-c1">3306</span>     <span class="pl-s1">to_run_exec</span>, <span class="pl-s1">to_run_interactive</span> <span class="pl-c1">=</span> [], <span class="pl-s1">nodelist</span>

<span class="pl-v">File</span> <span class="pl-c1">~</span><span class="pl-c1">/</span>.<span class="pl-s1">virtualenvs</span><span class="pl-c1">/</span><span class="pl-s1">science</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">10</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-v">IPython</span><span class="pl-c1">/</span><span class="pl-s1">core</span><span class="pl-c1">/</span><span class="pl-s1">interactiveshell</span>.<span class="pl-s1">py</span>:<span class="pl-c1">3394</span>, <span class="pl-s1">in</span> <span class="pl-v">InteractiveShell</span>.<span class="pl-s1">run_code</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
   <span class="pl-c1">3393</span> <span class="pl-s1">try</span>:
<span class="pl-c1">-&gt;</span> <span class="pl-c1">3394</span>     <span class="pl-k">if</span> <span class="pl-s1">async_</span>:
   <span class="pl-c1">3395</span>         <span class="pl-k">await</span> <span class="pl-en">eval</span>(<span class="pl-s1">code_obj</span>, <span class="pl-s1">self</span>.<span class="pl-s1">user_global_ns</span>, <span class="pl-s1">self</span>.<span class="pl-s1">user_ns</span>)

<span class="pl-v">Input</span> <span class="pl-v">In</span> [<span class="pl-c1">207</span>], <span class="pl-s1">in</span> <span class="pl-c1">&lt;</span><span class="pl-s1">cell</span> <span class="pl-s1">line</span>: <span class="pl-c1">10</span><span class="pl-c1">&gt;</span>()
      <span class="pl-c1">8</span>     <span class="pl-s1">return</span> <span class="pl-s1">x</span>
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">10</span> <span class="pl-s1">jax</span>.<span class="pl-en">linear_transpose</span>(<span class="pl-s1">foo</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">10</span>))(<span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">10</span>))

<span class="pl-v">Input</span> <span class="pl-v">In</span> [<span class="pl-c1">207</span>], <span class="pl-s1">in</span> <span class="pl-s1">foo</span>(<span class="pl-c1">**</span><span class="pl-c1">*</span><span class="pl-s1">failed</span> <span class="pl-s1">resolving</span> <span class="pl-s1">arguments</span><span class="pl-c1">**</span><span class="pl-c1">*</span>)
      <span class="pl-c1">1</span> <span class="pl-s1">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">p</span>):
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">2</span>     <span class="pl-s1">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">scan</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>, <span class="pl-s1">_</span>: (<span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-s1">p</span>, <span class="pl-c1">None</span>), <span class="pl-s1">jnp</span>.<span class="pl-en">zeros_like</span>(<span class="pl-s1">p</span>), <span class="pl-c1">None</span>, <span class="pl-s1">length</span><span class="pl-c1">=</span><span class="pl-c1">10</span>)[<span class="pl-c1">0</span>]

<span class="pl-v">JaxStackTraceBeforeTransformation</span>: <span class="pl-v">AssertionError</span>

<span class="pl-v">The</span> <span class="pl-s1">preceding</span> <span class="pl-s1">stack</span> <span class="pl-s1">trace</span> <span class="pl-c1">is</span> <span class="pl-s1">the</span> <span class="pl-s1">source</span> <span class="pl-s1">of</span> <span class="pl-s1">the</span> <span class="pl-v">JAX</span> <span class="pl-s1">operation</span> <span class="pl-s1">that</span>, <span class="pl-s1">once</span> <span class="pl-s1">transformed</span> <span class="pl-s1">by</span> <span class="pl-v">JAX</span>, <span class="pl-s1">triggered</span> <span class="pl-s1">the</span> <span class="pl-s1">following</span> <span class="pl-s1">exception</span>.

<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span>

<span class="pl-v">The</span> <span class="pl-s1">above</span> <span class="pl-s1">exception</span> <span class="pl-s1">was</span> <span class="pl-s1">the</span> <span class="pl-s1">direct</span> <span class="pl-s1">cause</span> <span class="pl-s1">of</span> <span class="pl-s1">the</span> <span class="pl-s1">following</span> <span class="pl-s1">exception</span>:

<span class="pl-v">AssertionError</span>                            <span class="pl-v">Traceback</span> (<span class="pl-s1">most</span> <span class="pl-s1">recent</span> <span class="pl-s1">call</span> <span class="pl-s1">last</span>)
<span class="pl-v">Input</span> <span class="pl-v">In</span> [<span class="pl-c1">207</span>], <span class="pl-s1">in</span> <span class="pl-c1">&lt;</span><span class="pl-s1">cell</span> <span class="pl-s1">line</span>: <span class="pl-c1">10</span><span class="pl-c1">&gt;</span>()
      <span class="pl-c1">7</span>         <span class="pl-s1">x</span> <span class="pl-c1">+=</span> <span class="pl-s1">p</span>
      <span class="pl-c1">8</span>     <span class="pl-s1">return</span> <span class="pl-s1">x</span>
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">10</span> <span class="pl-s1">jax</span>.<span class="pl-en">linear_transpose</span>(<span class="pl-s1">foo</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">10</span>))(<span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">10</span>))

    [... <span class="pl-s1">skipping</span> <span class="pl-s1">hidden</span> <span class="pl-c1">3</span> <span class="pl-s1">frame</span>]

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

<span class="pl-v">AssertionError</span>: </pre></div>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">v0.3.15</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">OSX ARM</p>
<p dir="auto"><em>first some motivation:</em></p>
<p dir="auto">I have a <strong>hermitian</strong> linear operator which is using a jvp followed by a vjp of complex (shallow) neural networks which ~ looks like this:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">matvec</span>(<span class="pl-s1">apply_fun</span>, <span class="pl-s1">params</span>, <span class="pl-s1">samples</span>, <span class="pl-s1">v</span>):
    <span class="pl-s1">_</span>, <span class="pl-s1">w</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">jvp</span>(<span class="pl-k">lambda</span> <span class="pl-s1">p</span>: <span class="pl-en">apply_fun</span>(<span class="pl-s1">p</span>, <span class="pl-s1">samples</span>), (<span class="pl-s1">params</span>,), (<span class="pl-s1">v</span>,))
    <span class="pl-s1">w</span> <span class="pl-c1">=</span> <span class="pl-s1">w</span>.<span class="pl-en">conj</span>()
    <span class="pl-s1">res</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">vjp</span>(<span class="pl-s1">apply_fun</span>, <span class="pl-s1">params</span>, <span class="pl-s1">samples</span>)[<span class="pl-c1">1</span>](<span class="pl-s1">w</span>)
    <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-en">tree_map</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-s1">conj</span>, <span class="pl-s1">res</span>)</pre></div>
<p dir="auto">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.</p>
<p dir="auto">I am using this linear operator with the cg and gmres iterative solvers, where its required to be linear_transpose'able.<br>
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.</p>
<p dir="auto">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:</p>
<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 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'))"><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">functools</span> <span class="pl-k">import</span> <span class="pl-s1">partial</span>

<span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-c1">100</span>
<span class="pl-s1">batchsize</span> <span class="pl-c1">=</span> <span class="pl-c1">5</span>
<span class="pl-s1">d</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>

<span class="pl-c"># generate a random jacobian and a vector</span>
<span class="pl-s1">k0</span>, <span class="pl-s1">k1</span>, <span class="pl-s1">k2</span>, <span class="pl-s1">k3</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">split</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">0</span>), <span class="pl-c1">4</span>)
<span class="pl-v">J</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">k0</span>, <span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-s1">n</span>,<span class="pl-s1">d</span>)) <span class="pl-c1">+</span> <span class="pl-c1">1j</span> <span class="pl-c1">*</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">k1</span>, <span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-s1">n</span>,<span class="pl-s1">d</span>))
<span class="pl-s1">v</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">k2</span>, <span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-s1">d</span>,)) <span class="pl-c1">+</span> <span class="pl-c1">1j</span> <span class="pl-c1">*</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">k3</span>, <span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-s1">d</span>,))

<span class="pl-c"># the linar operator</span>
<span class="pl-k">def</span> <span class="pl-en">mv</span>(<span class="pl-v">J</span>, <span class="pl-s1">v</span>):
    <span class="pl-k">return</span> <span class="pl-v">J</span>.<span class="pl-v">T</span>.<span class="pl-en">conj</span>()@(<span class="pl-v">J</span>@<span class="pl-s1">v</span>)

<span class="pl-c"># works fine</span>
<span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">scipy</span>.<span class="pl-s1">sparse</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">cg</span>(<span class="pl-en">partial</span>(<span class="pl-s1">mv</span>, <span class="pl-v">J</span>), <span class="pl-s1">v</span>)
<span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">scipy</span>.<span class="pl-s1">sparse</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">gmres</span>(<span class="pl-en">partial</span>(<span class="pl-s1">mv</span>, <span class="pl-v">J</span>), <span class="pl-s1">v</span>, <span class="pl-s1">solve_method</span><span class="pl-c1">=</span><span class="pl-s">'incremental'</span>)

<span class="pl-c"># mv is hermitian, so its transpose is trivial</span>
<span class="pl-k">def</span> <span class="pl-en">mv_transposed</span>(<span class="pl-v">J</span>, <span class="pl-s1">v</span>):
    <span class="pl-k">return</span> <span class="pl-en">mv</span>(<span class="pl-v">J</span>, <span class="pl-s1">v</span>.<span class="pl-en">conj</span>()).<span class="pl-en">conj</span>()


<span class="pl-c"># now the operator where the vjp is looping over small batches</span>
<span class="pl-k">def</span> <span class="pl-en">mv_batched</span>(<span class="pl-v">J</span>, <span class="pl-s1">v</span>):    
    <span class="pl-s1">w</span> <span class="pl-c1">=</span> <span class="pl-v">J</span>@<span class="pl-s1">v</span>

    <span class="pl-v">J_batched</span> <span class="pl-c1">=</span> <span class="pl-v">J</span>.<span class="pl-en">reshape</span>((<span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">batchsize</span>)<span class="pl-c1">+</span><span class="pl-v">J</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">1</span>:])
    <span class="pl-s1">w_batched</span> <span class="pl-c1">=</span> <span class="pl-s1">w</span>.<span class="pl-en">reshape</span>((<span class="pl-c1">-</span><span class="pl-c1">1</span>, <span class="pl-s1">batchsize</span>)<span class="pl-c1">+</span><span class="pl-s1">w</span>.<span class="pl-s1">shape</span>[<span class="pl-c1">1</span>:])
    <span class="pl-v">Jw</span> <span class="pl-c1">=</span> <span class="pl-v">J_batched</span>, <span class="pl-s1">w_batched</span>
    
    <span class="pl-k">def</span> <span class="pl-en">_mv</span>(<span class="pl-v">Jw</span>):
        <span class="pl-v">J</span>, <span class="pl-s1">w</span> <span class="pl-c1">=</span> <span class="pl-v">Jw</span>
        <span class="pl-k">return</span> (<span class="pl-s1">w</span>.<span class="pl-en">conj</span>() @ <span class="pl-v">J</span>).<span class="pl-en">conj</span>() 
    
    <span class="pl-k">def</span> <span class="pl-en">f</span>(<span class="pl-s1">carry</span>, <span class="pl-s1">x</span>):
        <span class="pl-k">return</span> <span class="pl-s1">carry</span> <span class="pl-c1">+</span> <span class="pl-en">_mv</span>(<span class="pl-s1">x</span>), <span class="pl-c1">None</span>
    
    <span class="pl-s1">res</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">scan</span>(<span class="pl-s1">f</span>, <span class="pl-s1">jnp</span>.<span class="pl-en">zeros_like</span>(<span class="pl-s1">v</span>), <span class="pl-v">Jw</span>, <span class="pl-s1">unroll</span><span class="pl-c1">=</span><span class="pl-c1">1</span>)
    
    <span class="pl-k">return</span> <span class="pl-s1">res</span>


<span class="pl-k">assert</span> <span class="pl-s1">jnp</span>.<span class="pl-en">allclose</span>(<span class="pl-en">mv</span>(<span class="pl-v">J</span>,<span class="pl-s1">v</span>), <span class="pl-en">mv_batched</span>(<span class="pl-v">J</span>,<span class="pl-s1">v</span>))
<span class="pl-k">assert</span> <span class="pl-s1">jnp</span>.<span class="pl-en">allclose</span>(<span class="pl-en">mv_transposed</span>(<span class="pl-v">J</span>,<span class="pl-s1">v</span>), <span class="pl-s1">jax</span>.<span class="pl-en">linear_transpose</span>(<span class="pl-en">partial</span>(<span class="pl-s1">mv</span>, <span class="pl-v">J</span>), <span class="pl-s1">v</span>)(<span class="pl-s1">v</span>)[<span class="pl-c1">0</span>])

<span class="pl-c"># linear_transposing mv_batched gives the error:</span>
<span class="pl-en">print</span>( <span class="pl-s1">jax</span>.<span class="pl-en">linear_transpose</span>(<span class="pl-en">partial</span>(<span class="pl-s1">mv_batched</span>, <span class="pl-v">J</span>), <span class="pl-s1">v</span>)(<span class="pl-s1">v</span>)[<span class="pl-c1">0</span>] )

<span class="pl-c"># goal: use mv_batched with the iterative solvers</span>
<span class="pl-c"># same error:</span>
<span class="pl-c"># print( jax.scipy.sparse.linalg.cg(partial(mv_batched, J), v))</span>
<span class="pl-c"># print( jax.scipy.sparse.linalg.gmres(partial(mv_batched, J), v, solve_method='incremental'))</span></pre></div>
<p dir="auto">And the output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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)
&lt;ipython-input-1-dbeaca6d614d&gt; in &lt;module&gt;
     45 # jax.scipy.sparse.linalg.gmres(partial(mv_batched, J), v, solve_method='incremental')
     46 
---&gt; 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,
-&gt; 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:
--&gt; 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)
-&gt; 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: "><pre class="notranslate"><code class="notranslate">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)
&lt;ipython-input-1-dbeaca6d614d&gt; in &lt;module&gt;
     45 # jax.scipy.sparse.linalg.gmres(partial(mv_batched, J), v, solve_method='incremental')
     46 
---&gt; 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,
-&gt; 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:
--&gt; 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)
-&gt; 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: 
</code></pre></div>
<p dir="auto">[I thought of just using linear_call, however it's still missing a XLA translation rule so it currently can't be jitted]</p>
<p dir="auto">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).</p>