Fix typo in Pallas design

PiperOrigin-RevId: 621275025

Author: sharadmv

docs/pallas/design.md

@@ -236,10 +236,11 @@ add = pl.pallas_call(
     add_kernel,
     out_shape=jax.ShapeDtypeStruct((8,), jnp.int32),
     in_specs=[
-      pl.BlockSpec(lambda i:, i, (2,)),
-      pl.BlockSpec(lambda i:, i, (2,))
+      pl.BlockSpec(lambda i: i, (2,)),
+      pl.BlockSpec(lambda i: i, (2,))
     ],
-    out_specs=pl.BlockSpec(lambda i: i, (2,))
+    out_specs=pl.BlockSpec(lambda i: i, (2,)),
+    grid=(4,))
 add(x, y)
 ```