[Mosaic GPU] Fix some typos in docs

Author: justinjfu

docs/pallas/gpu/reference.md

@@ -700,9 +700,9 @@ def run_kernel(x_ref, y_ref):
   def kernel_body():
     # Once we enter the pl.core_map scope, we are in the body of the kernel.
     block_slice = pl.ds(lax.axis_index("x") * 128, 128)
-    o_ref[block_slice] = x_ref[block_slice] + 1
+    y_ref[block_slice] = x_ref[block_slice] + 1
 
-x = jnp.arange(128, jnp.float32)
+x = jnp.arange(256, jnp.float32)
 y_init = jnp.zeros_like(x)
 y = run_kernel(x, y_init)
 np.testing.assert_array_equal(y, x + 1)
@@ -721,12 +721,12 @@ mesh = plgpu.Mesh(grid=(2,), grid_names=("x",))
     out_shape=jax.ShapeDtypeStruct((256,), jnp.float32),
     mesh=mesh
 )
-def increment_kernel_core_map(x_ref, y_ref):
+def run_kernel(x_ref, y_ref):
   # x_ref and y_ref are in GMEM!
   block_slice = pl.ds(lax.axis_index("x") * 128, 128)
-  o_ref[block_slice] = x_ref[block_slice] + 1
+  y_ref[block_slice] = x_ref[block_slice] + 1
 
-x = jnp.arange(128, jnp.float32)
+x = jnp.arange(256, jnp.float32)
 y = run_kernel(x)  # No need to preallocate outputs as in pl.core_map.
 np.testing.assert_array_equal(y, x + 1)
 ```
@@ -760,7 +760,7 @@ def run_kernel(x_ref, y_ref, barrier_ref):
 
   @pl.when(thread_id == 0)
   def producer_thread():
-    smem_val = x_ref[...] + 1
+    x_ref[...] = x_ref[...] + 1
     plgpu.barrier_arrive(barrier_ref)  # Signal the consumer thread
 
   @pl.when(thread_id == 1)