Propose to improve p01 and p02

book/src/puzzle_02/puzzle_02.md

@@ -130,4 +130,4 @@ While this direct indexing works for simple element-wise operations, consider:
 - What if we need to broadcast one array to another?
 - How to ensure coalesced access across multiple arrays?
 
-These questions will be addressed when we [introduce LayoutTensor in Puzzle 4](../puzzle_04/).
+These questions will be addressed when we [introduce LayoutTensor in Puzzle 4](../puzzle_04/introduction_layout_tensor.md).

problems/p01/p01.mojo

@@ -1,3 +1,4 @@
+from memory import UnsafePointer
 from gpu import thread_idx
 from gpu.host import DeviceContext
 from testing import assert_equal

problems/p02/p02.mojo

@@ -1,5 +1,5 @@
 from memory import UnsafePointer
-from gpu import thread_idx, block_dim, block_idx
+from gpu import thread_idx
 from gpu.host import DeviceContext
 from testing import assert_equal
 

solutions/p02/p02.mojo

@@ -1,5 +1,5 @@
 from memory import UnsafePointer
-from gpu import thread_idx, block_dim, block_idx
+from gpu import thread_idx
 from gpu.host import DeviceContext
 from testing import assert_equal