jax.lax.stop_gradient should be not useless in an example

The `x` in `u(x)` is different from one in lambdas, so `stop_gradient` was doing nothing.
I have fixed the example for it to have some effect.

Author: faucct

chapter_preliminaries/autograd.md

@@ -463,10 +463,10 @@ import jax
 
 y = lambda x: x * x
 # jax.lax primitives are Python wrappers around XLA operations
-u = jax.lax.stop_gradient(y(x))
-z = lambda x: u * x
-
-grad(lambda x: z(x).sum())(x) == y(x)
+(
+    grad(lambda x: (y(x) * x).sum())(x) == 3 * y(x),
+    grad(lambda x: (jax.lax.stop_gradient(y(x)) * x).sum())(x) == y(x),
+)
 ```
 
 Note that while this procedure
@@ -694,4 +694,4 @@ For now, try to remember these basics: (i) attach gradients to those variables w
 
 :begin_tab:`jax`
 [Discussions](https://discuss.d2l.ai/t/17970)
-:end_tab:
+:end_tab: