Repository: google/jax

Issue:
`lax.scan` is ~6x slower to run than hand-written loops
This is on a GPU backend, haven't tried on others.

https://colab.research.google.com/drive/1N1jrvGNFRhTnYLOCL6vXUzJN3pRR3pdr

Minimal repro below
```
from jax import numpy as np
from jax import grad, jit, vmap, lax
from jax import random as jax_random

import numpy as onp

@jit
def rewards_to_go(rewards, mask, gamma=0.99):
  r"""Computes rewards to go.
  Args:
    rewards: np.ndarray of shape (B, T) of rewards.
    mask: np.ndarray of shape (B, T) of mask for the rewards.
    gamma: float, discount factor.

  Returns:
    rewards to go, np.ndarray of shape (B, T).
  """
  B, T = rewards.shape  # pylint: disable=invalid-name,unused-variable

  masked_rewards = rewards * mask  # (B, T)

  # Compute r2g_{T-1} at the start and then compute backwards in time.
  r2gs = [masked_rewards[:, -1]]

  # Go from T-2 down to 0.
  for t in reversed(range(T - 1)):
    r2gs.append(masked_rewards[:, t] + (gamma * r2gs[-1]))

  # The list should have length T.
  assert T == len(r2gs)

  # First we stack them in the correct way to make it (B, T), but these are
  # still from newest (T-1) to oldest (0), so then we flip it on time axis.
  return np.flip(np.stack(r2gs, axis=1), axis=1)


@jit
def scan_rewards_to_go(rewards, mask, gamma=0.99):
  masked_rewards = rewards * mask  # (B, T)

  reversed_rewards = np.flip(masked_rewards, axis=1)  # (B, T) flipped on time.
  rrt = np.transpose(reversed_rewards)  # (T, B) transpose to scan over time.

  def discounting_add(carry, reward):
    x = reward + (gamma * carry)
    return x, x

  _, ys = lax.scan(discounting_add,
                   np.zeros_like(rrt[0], dtype=np.float32),
                   rrt.astype(np.float32))

  # ys is (T, B) and T is in reverse order.
  return np.flip(np.transpose(ys), axis=1)

B, T = 16, 128

num_examples = 100
rewards = []
pvs = []
mask = []

for _ in range(num_examples):
  rewards.append(onp.random.randn(B, T))
  pvs.append(onp.random.randn(B, T+1))
  ones = onp.full((B, T), 1, dtype=onp.int32)
  for one in ones:
    l = onp.random.randint(0, T)
    one[range(l,T)] = 0
  mask.append(ones)
```

Now time the invocations:

```
%timeit [rewards_to_go(rewards[i], mask[i]) for i in range(num_examples)]
```

and 

```
%timeit [scan_rewards_to_go(rewards[i], mask[i]) for i in range(num_examples)]
```


Discussion:
I have a related scan file here: https://gist.github.com/cshesse/73f27d42c356e6e763275432422dacbe

On CPU:
```
reward_to_go_jax 0.028926010942086577
reward_to_go_scipy 0.001616921043023467
scipy speedup 17.889563047554923
```

In the CPU case, for large values of N, the performance gets much closer to scipy lfilter.

On GPU:
```
reward_to_go_jax 5.656240948010236
reward_to_go_scipy 0.001609566854313016
scipy speedup 3514.1385602304745
```

`scipy.signal.lfilter` may not be a useful benchmark to compare against, but the jax version seems to get significantly slower on GPU.

This is using `git+https://github.com/google/jax.git@11c512a194add25ade20dbfacf635d6e0834eba3` and `https://storage.googleapis.com/jax-wheels/cuda92/jaxlib-0.1.15-cp37-none-linux_x86_64.whl`

@jlebar Any gut reactions here?

I've been meaning to dump this JAX code (both in the OP and in @cshesse's gist) into HLO text for easier repros, but I haven't gotten to it yet. Let me know if you can take a look without that step :)

Wild guess: could these slowdowns be related to the TupleThunk issue you just fixed?
It looks like this benchmark is including compilation time?

> Wild guess: could these slowdowns be related to the TupleThunk issue you just fixed?

Meaning, the TupleThunk change makes things slower?  That would be surprising.  Possible, I suppose, but I don't immediately see how.

Other than the benchmarking methodology question I don't have an immediate guess here.  I don't know enough about the difference in HLO implementation of the two functions to have any ideas, without actually looking at the HLO.

I'm always trying to think bigger picture, though. We're getting a fair number of these performance bugs, which is genuinely great.  And the bugs are also fairly high quality.  But how do we streamline the process of receiving them, triaging them, and acting upon them, for you and for us?

I'm wondering if e.g. we could have a script/document that people can follow when they're reporting a performance bug, which will capture the textual HLO and xla_hlo_profile?  Ideally it could capture some sort of execution trace too, but I'm not sure how much of that is OSS at the moment.

Thinking even bigger, if JAX had some sort of benchmark framework then it could take care of ensuring that we ignore the compilation time, in addition to doing all of the things above for bug reports.
> Meaning, the TupleThunk change makes things slower? That would be surprising. Possible, I suppose, but I don't immediately see how.

Oh, sorry, I meant the opposite but my wording was bad: these numbers are all from before the TupleThunk fix, so I was guessing that your TupleThunk fix could have improved things.

Great thoughts on the big picture. We should at least have a doc, and I bet we can make a script, including maybe an internal one that collects a trace (which we can run ourselves on small repros). Maybe we (on the JAX side) can take this opportunity to start developing those things!

Thanks for weighing in, even with very little info from our side here.
In my script, I call the function once during setup, is that not enough to avoid capturing compilation time?

For a benchmarking script, maybe a `jax.experimental.benchmark(fn)` vaguely like go bench:

```python
import time
import jaxlib

first_run = True


def benchmark(fn, duration_sec=10):
    global first_run
    if first_run:
        print("jax version", jax.__version__)
        print("jaxlib version", jaxlib.__version__)
    first_run = False
    fn(1)
    n = 1
    while True:
        start = time.time()
        fn(n)
        elapsed = time.time() - start
        if elapsed > duration_sec:
            break
        n *= 10
    print("%s %d calls  %f us/call" % (fn.__name__, n, elapsed / n * 1e6))


def benchmark_reward_to_go_scipy(n):
    for _ in range(n):
        reward_to_go_scipy(rews, 0.99)


def benchmark_reward_to_go_jax(n):
    for _ in range(n):
        reward_to_go_jax(jrews, 0.99).block_until_ready()


benchmark(benchmark_reward_to_go_scipy)
benchmark(benchmark_reward_to_go_jax)
```
> In my script, I call the function once during setup, is that not enough
to avoid capturing compilation time?

Should be, yes.

On Thu, Jun 6, 2019 at 9:30 AM cshesse <notifications@github.com> wrote:

> In my script, I call the function once during setup, is that not enough to
> avoid capturing compilation time?
>
> For a benchmarking script, maybe a jax.experimental.benchmark(fn) vaguely
> like go bench setup:
>
> import time
> import jaxlib
>
> first_run = True
>
>
> def benchmark(fn, duration_sec=10):
>     global first_run
>     if first_run:
>         print("jax version", jax.__version__)
>         print("jaxlib version", jaxlib.__version__)
>     first_run = False
>     fn(1)
>     n = 1
>     while True:
>         start = time.time()
>         fn(n)
>         elapsed = time.time() - start
>         if elapsed > duration_sec:
>             break
>         n *= 10
>     print("%s %d calls  %f us/call" % (fn.__name__, n, elapsed / n * 1e6))
>
>
> def benchmark_reward_to_go_scipy(n):
>     for _ in range(n):
>         reward_to_go_scipy(rews, 0.99)
>
>
> def benchmark_reward_to_go_jax(n):
>     for _ in range(n):
>         reward_to_go_jax(jrews, 0.99)
>
>
> benchmark(benchmark_reward_to_go_scipy)
> benchmark(benchmark_reward_to_go_jax)
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/google/jax/issues/810?email_source=notifications&email_token=AABEZB5LLNWV7NJYD55MXCDPZE3QVA5CNFSM4HSJ4MHKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODXDNLWY#issuecomment-499570139>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AABEZBYK2I4BXXAUVG7CF7DPZE3QVANCNFSM4HSJ4MHA>
> .
>

The main problem with the example in the first comment is that there is almost no useful work in the loop (a elementwise multiply and an add of `f32[16]` vector). You won't get good GPU utilization from this no matter what we do; CPU will be faster. The loop benchmark (using `scan`) is mostly measuring loop overheads; we could probably lower them but I am wondering if this code is actually representative of something you want to do. Put another way: even if the `scan` variant is 6x slower, does it actually matter?


@hawkinsp I've ran into this with bigger workloads. Specifically, a scan implementation lead to about a ~4x slow down, ~6x larger GPU memory usage, ~4-6x higher transfer rates as reported by `nvidia-smi` than an unrolled implementation. I'm not so certain that this is just an issue of higher overhead.

For a bit of context, the workload at each iterations was:
1. `extract_image_patch(Nx28x28x2)`
2. `tensordot(Nx28x28x9, Nx28x28x9x10 -> Nx28x28x10)`
3. `max(Nx28x28x10, axis=-1)`
with batch size `N=10`. While the scan version compiles very quickly, the slow down makes it completely impractical. This was on an Nvidia 2080 Ti.
Additional observations: for the workload I describe earlier, if I run 25 unrolled steps per scan iteration, I still see a drastic increase in GPU usage, memory usage and CPU-GPU communication. Everything I've observed seems to hint at some memory reuse issue or some arrays placed on the wrong device issue. This issue is a big blocker for me. If there is anything I can do to help fix it, please let me know.

Note that while the reported GPU usage increases (from ~60% to 99%) when going from unrolled to scan, the actual performance is much much worse.
@gehring Would it be possible to get us a self-contained example that demonstrates your slowdown? There are a few known drawbacks of XLA:GPU loops, but `scan` is typically supposed to be able to avoid the worst of it (raw `while_loop` loops have to bring their condition boolean back to the host at each step, while `scan` always has a fixed trip count).
@jekbradbury I'll try to strap a standalone example by the end of the week. In the meantime, if it makes any difference, I realize I forgot to mention that the workload I described is part of a SGD (pure python) loop, doing a forward and backwards pass in each jitted call.
@jekbradbury It ended up being much easier to do than I thought. The code where I first encountered this is somewhat complex but I was able to reproduce it with just a scan looping `amat @ x + bvec`. I saved it in a [notebook](https://colab.research.google.com/gist/gehring/bee2f5bd37c108c7620404e2c47652ed/jax-scan-issue.ipynb) for convenience.

On Google's public colab (GPU runtime), I get the following results:
```
Time taken unrolled: 12.955157 sec
Time taken scan: 197.164844 sec
```
I tried this example on a different workstation (with an old Quadro K2200) and observed pretty much the same symptoms as described earlier (using `nvidia-smi`).
Thanks! This looks like a performance issue with JVP and possibly transpose of scan, which emit `dynamic-slice` and `dynamic-update-slice` operations (respectively) that appear to be the bottleneck here (and seemingly not necessary, since they don't show up in the unrolled version). Going to look into it, but I don't know this part of the code very well.
@jekbradbury Any progress or updates on this? No pressure, I'm just trying to get an idea as to whether a fix is in the works or if I should commit to a workaround.
I had a partial solution in the linked PR, but we ended up deciding against that approach and sketched out a more direct (and complete) solution two weeks ago. I haven’t implemented it yet, though it’s been on my todos; I expect I’ll have a new PR sometime in the next week or so.

Resolving patch:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -617,34 +617,81 @@ def _scan_partial_eval(trace, *tracers, **kwargs):
       carry_uk = _map(operator.or_, carry_uk, carry_uk_out)
   else:
     assert False, "Fixpoint not reached"
+  num_res = len(jaxpr_1.out_avals) - len(jaxpr_2.out_avals)
+
+  # The residuals are treated as extensive outputs of jaxpr_1 (and extensive
+  # inputs to jaxpr_2), but residuals that are loop-invariant can be hoisted.
+  # TODO(mattjj): hoist other loop-invariant values here too (instantiate=False)
+  invariant_pvals = [pe.PartialVal((None, core.unit if uk else t.pval[1]))
+                     for uk, t in zip(unknowns[:num_consts], tracers[:num_consts])]
+  other_pvals = [pe.PartialVal((a, core.unit)) for a in jaxpr_1.in_avals[num_consts:]]
+  in_pvals_1 = invariant_pvals + other_pvals
+  untyped_jaxpr_1, out_pvals_1, consts_1 = pe.trace_to_jaxpr(
+      lu.wrap_init(core.jaxpr_as_fun(jaxpr_1)), in_pvals_1,
+      instantiate=[True] * (num_carry + num_ys) + [False] * num_res)
+  const_avals_1 = [raise_to_shaped(core.get_aval(c)) for c in consts_1]
+  in_avals_1 = [core.abstract_unit] * num_consts + jaxpr_1.in_avals[num_consts:]
+  out_avals_1 = [core.abstract_unit if pv is None else pv for pv, c in out_pvals_1]
+  jaxpr_1_opt = pe.TypedJaxpr(pe.closure_convert_jaxpr(untyped_jaxpr_1),
+                              (), const_avals_1 + in_avals_1, out_avals_1)
+  num_consts_1 = num_consts + len(consts_1)
+  # any now-known residuals are intensive, so we want to revise jaxpr_2 to take
+  # those inputs as constants rather than as extensive inputs
+  _, _, res_pvals = split_list(out_pvals_1, [num_carry, num_ys])
+  intensive_residuals = [const for pv, const in res_pvals if pv is None]
+  move = [False] * len(jaxpr_1.in_avals) + [pv is None for pv, _ in res_pvals]
+  jaxpr_2_opt = _move_binders_to_front(jaxpr_2, move)
+  num_consts_2 = num_consts + len(intensive_residuals)
+
+  in_consts = (list(consts_1) + [core.unit] * num_consts +
+               [core.unit if uk else t.pval[1]
+                for uk, t in zip(unknowns[num_consts:], tracers[num_consts:])])
+  linear_1 = ([False] * len(consts_1) + [True] * num_consts +
+              [lin or uk for uk, lin in zip(unknowns[num_consts:], linear[num_consts:])])
+  out_flat = scan_p.bind(
+      *in_consts, forward=forward, length=length, jaxpr=jaxpr_1_opt,
+      num_consts=num_consts_1, num_carry=num_carry, linear=linear_1)
+  out_carry, ys, res_and_units = split_list(out_flat, [num_carry, num_ys])
+  extensive_residuals = [r for r, (pv, _) in zip(res_and_units, res_pvals) if pv is not None]
 
-  in_consts = [core.unit if uk else t.pval[1] for uk, t in zip(unknowns, tracers)]
   new_tracers = [trace.instantiate_const(t) if uk else trace.new_instantiated_literal(core.unit)
                  for uk, t in zip(unknowns, tracers)]
-
   carry_avals, y_avals = split_list(jaxpr.out_avals, [num_carry])
   ys_avals = _map(partial(_promote_aval_rank, length), y_avals)
   out_avals = carry_avals + ys_avals
   out_pvs = [aval if uk else None for aval, uk in zip(out_avals, out_uk)]
 
-  linear_1 = [lin or uk for uk, lin in zip(unknowns, linear)]
-  out_flat = scan_p.bind(
-      *in_consts, forward=forward, length=length, jaxpr=jaxpr_1,
-      num_consts=num_consts, num_carry=num_carry, linear=linear_1)
-  out_carry, ys, residuals = split_list(out_flat, [num_carry, num_ys])
   out_consts = out_carry + ys
-  residual_tracers = _map(trace.new_instantiated_const, residuals)
+  int_res_tracers = _map(trace.new_instantiated_const, intensive_residuals)
+  ext_res_tracers = _map(trace.new_instantiated_const, extensive_residuals)
   out_tracers = [pe.JaxprTracer(trace, pe.PartialVal((pv, const)), None)
                  for pv, const in zip(out_pvs, out_consts)]
-  linear_2 = ([lin or not uk for uk, lin in zip(unknowns, linear)]
-              + [False] * len(residual_tracers))
-  eqn = pe.new_jaxpr_eqn(new_tracers + residual_tracers, out_tracers, scan_p,
-                         (), dict(forward=forward, length=length, jaxpr=jaxpr_2,
-                                  num_consts=num_consts, num_carry=num_carry,
-                                  linear=linear_2))
+  linear_2 = ([False] * len(int_res_tracers) +
+              [lin or not uk for uk, lin in zip(unknowns, linear)] +
+              [False] * len(ext_res_tracers))
+  eqn = pe.new_jaxpr_eqn(int_res_tracers + new_tracers + ext_res_tracers,
+                         out_tracers, scan_p, (),
+                         dict(forward=forward, length=length, jaxpr=jaxpr_2_opt,
+                              num_consts=num_consts_2,
+                              num_carry=num_carry, linear=linear_2))
   for t in out_tracers: t.recipe = eqn
   return out_tracers
 
+def _move_binders_to_front(typed_jaxpr, to_move):
+  assert not typed_jaxpr.jaxpr.constvars and not typed_jaxpr.jaxpr.freevars
+  assert len(typed_jaxpr.in_avals) == len(to_move)
+  new_invars = _move_to_front(typed_jaxpr.jaxpr.invars, to_move)
+  new_jaxpr = core.Jaxpr((), (), new_invars, typed_jaxpr.jaxpr.outvars,
+                         typed_jaxpr.jaxpr.eqns)
+  new_in_avals = _move_to_front(typed_jaxpr.in_avals, to_move)
+  new_typed_jaxpr = core.TypedJaxpr(new_jaxpr, typed_jaxpr.literals,
+                                    new_in_avals, typed_jaxpr.out_avals)
+  return new_typed_jaxpr
+
+def _move_to_front(lst, to_move):
+  return ([elt for elt, move in zip(lst, to_move) if move] +
+          [elt for elt, move in zip(lst, to_move) if not move])
+
 def _promote_aval_rank(sz, aval):
   if aval is core.abstract_unit:
     return core.abstract_unit
@@ -655,54 +702,66 @@ def _scan_transpose(cts, *args, **kwargs):
   forward, length, num_consts, num_carry, jaxpr, linear = split_dict(
       kwargs, ["forward", "length", "num_consts", "num_carry", "jaxpr", "linear"])
 
-  # we can only transpose scans for which the nonlinear values appear in xs
+  # we've only implemented transposing scans with specific lin/nonlin patterns
   consts_lin, init_lin, xs_lin = split_list(linear, [num_consts, num_carry])
-  num_lin = sum(xs_lin)
-  if not all(consts_lin) or not all(init_lin) or not all(xs_lin[:num_lin]):
+  num_ires = len(consts_lin) - sum(consts_lin)
+  num_eres = len(xs_lin) - sum(xs_lin)
+  if consts_lin != [False] * num_ires + [True] * (len(consts_lin) - num_ires):
+    raise NotImplementedError
+  if xs_lin != [True] * (len(xs_lin) - num_eres) + [False] * num_eres:
+    raise NotImplementedError
+  if not all(init_lin):
     raise NotImplementedError
 
-  consts, init, xs, res = split_list(args, [num_consts, num_carry, num_lin])
-  assert not any(r is ad.undefined_primal for r in res)
+  consts, init, xs = split_list(args, [num_consts, num_carry])
+  ires, consts = split_list(consts, [num_ires])
+  xs, eres = split_list(xs, [sum(xs_lin)])
+  assert not any(r is ad.undefined_primal for r in ires)
+  assert not any(r is ad.undefined_primal for r in eres)
 
   carry_avals, y_avals = split_list(jaxpr.out_avals, [num_carry])
   ys_avals = _map(partial(_promote_aval_rank, length), y_avals)
   ct_carry, ct_ys = split_list(cts, [num_carry])
   ct_carry = _map(ad.instantiate_zeros_aval, carry_avals, ct_carry)
   ct_ys = _map(ad.instantiate_zeros_aval, ys_avals, ct_ys)
-  ct_consts = _map(ad_util.zeros_like_aval, jaxpr.in_avals[:num_consts])
+  ct_consts = _map(ad_util.zeros_like_aval, jaxpr.in_avals[num_ires:num_consts])
 
-  #       jaxpr :: [T d] -> [T c] -> [T a, res] -> ([T c], [T b])
-  # jaxpr_trans :: [] -> [CT d, CT c] -> [CT b, res] -> ([CT d, CT c], [CT a])
-  jaxpr_trans = _transpose_jaxpr(num_consts, len(res), jaxpr)
-  linear_trans = ([True] * (len(ct_consts) + len(ct_carry) + len(ct_ys))
-                  + [False] * len(res))
+  #       jaxpr :: [ires, T d] -> [T c] -> [T a, eres] -> ([T c], [T b])
+  # jaxpr_trans :: [ires] -> [CT d, CT c] -> [CT b, eres] -> ([CT d, CT c], [CT a])
+  jaxpr_trans = _transpose_jaxpr(num_ires, num_consts - num_ires, num_eres, jaxpr)
+  linear_trans = ([False] * num_ires +
+                  [True] * (len(ct_consts) + len(ct_carry) + len(ct_ys)) +
+                  [False] * num_eres)
 
   outs = scan_p.bind(
-      *(ct_consts + ct_carry + ct_ys + res), forward=not forward, length=length,
-      jaxpr=jaxpr_trans, num_consts=0, num_carry=num_consts+num_carry,
-      linear=linear_trans)
-  ct_consts, ct_init, ct_xs = split_list(outs, [num_consts, num_carry])
-  return ct_consts + ct_init + ct_xs + [None] * len(res)
-
-# transpose_jaxpr :: ([c, a, res] -> b) -> ([CT c, CT b, res] -> [CT c, CT a]
-def _transpose_jaxpr(num_c, num_res, jaxpr):
-  num_a = len(jaxpr.in_avals) - num_c - num_res
-  c_avals, a_avals, res_avals = split_list(jaxpr.in_avals, [num_c, num_a])
+      *(ires + ct_consts + ct_carry + ct_ys + eres), forward=not forward,
+      length=length, jaxpr=jaxpr_trans, num_consts=num_ires,
+      num_carry=num_consts-num_ires+num_carry, linear=linear_trans)
+  ct_consts, ct_init, ct_xs = split_list(outs, [num_consts - num_ires, num_carry])
+  return [None] * num_ires + ct_consts + ct_init + ct_xs + [None] * num_eres
+
+# transpose_jaxpr :: ([res1, c, a, res2] -> b)
+#                    -> ([res1, CT c, CT b, res2] -> [CT c, CT a])
+def _transpose_jaxpr(num_res1, num_c, num_res2, jaxpr):
+  num_a = len(jaxpr.in_avals) - num_res1 - num_c - num_res2
+  res1_avals, c_avals, a_avals, res2_avals = split_list(
+      jaxpr.in_avals, [num_res1, num_c, num_a])
   num_b = len(jaxpr.out_avals)
   b_avals = list(jaxpr.out_avals)
 
   @lu.wrap_init
-  def transposed(*cbar_bbar_res):
-    c_bar, b_bar, res = split_list(cbar_bbar_res, [num_c, num_b])
-    primals = [ad.undefined_primal] * (num_c + num_a) + res
+  def transposed(*res1_cbar_bbar_res2):
+    res1, c_bar, b_bar, res2 = split_list(
+        res1_cbar_bbar_res2, [num_res1, num_c, num_b])
+    primals = res1 + [ad.undefined_primal] * (num_c + num_a) + res2
     _, cbar_abar = ad.backward_pass(jaxpr.jaxpr, jaxpr.literals, (), primals,
                                     b_bar)
-    new_c_bar, a_bar, _ = split_list(cbar_abar, [num_c, num_a])
+    _, new_c_bar, a_bar, _ = split_list(cbar_abar, [num_res1, num_c, num_a])
     a_bar = _map(ad.instantiate_zeros_aval, a_avals, a_bar)
     c_bar = _map(ad.instantiate_zeros_aval, c_avals,
                 _map(ad.add_tangents, c_bar, new_c_bar))
     return c_bar + a_bar
-  return _make_typed_jaxpr(transposed, c_avals + b_avals + res_avals)
+  return _make_typed_jaxpr(transposed, res1_avals + c_avals + b_avals + res2_avals)
 
 def _make_typed_jaxpr(traceable, in_avals):
   pvals = [pe.PartialVal((aval, core.unit)) for aval in in_avals]


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -1365,6 +1365,20 @@ def bad_matvec_usage(a):
     with self.assertRaisesRegex(ValueError, re.escape("matvec() output shapes")):
       api.jvp(bad_matvec_usage, (1.0,), (1.0,))
 
+  def testIssue810(self):
+    def loss(A):
+      def step(x, i):
+        return np.matmul(A, x), None
+      init_x = np.zeros(A.shape[-1:])
+      last_x, _ = lax.scan(step, init_x, np.arange(10))
+      return np.sum(last_x)
+
+    A = np.zeros((3, 3))
+    # The second DUS was unnecessarily replicating A across time.
+    # We check XLA because _scan_impl is "underneath" the jaxpr language.
+    s = str(api.xla_computation(api.grad(loss))(A).GetHloText())
+    assert s.count("dynamic-update-slice(") < 2
+
 
 if __name__ == '__main__':
   absltest.main()