diff --git a/tests/mosaic/gpu_layout_inference_test.py b/tests/mosaic/gpu_layout_inference_test.py
index b0c6915d8536..7c477326c989 100644
--- a/tests/mosaic/gpu_layout_inference_test.py
+++ b/tests/mosaic/gpu_layout_inference_test.py
@@ -413,7 +413,9 @@ def test_infer_layout_from_body_op_to_yield_op_to_for_op(self):
     shape = (64, 64)
     with ir.InsertionPoint(self.module.body):
       c_ty = ir.VectorType.get(shape, ir.BF16Type.get())
-      ab_type = ir.MemRefType.get(shape, ir.BF16Type.get())
+      ab_type = ir.MemRefType.get(
+          shape, ir.BF16Type.get(), memory_space=mgpu.utils.smem()
+      )
       i32 = ir.IntegerType.get_signless(32)
       lower_bound, upper_bound, step, a, b, c = undefs(
           i32, i32, i32, ab_type, ab_type, c_ty
@@ -794,7 +796,7 @@ def test_infer_wgmma_layout_correctly(self, lhs_memory_space):
 
     with ir.InsertionPoint(self.module.body):
       vec_ty = ir.VectorType.get(shape, f32)
-      ref_ty = ir.MemRefType.get(shape, f32)
+      ref_ty = ir.MemRefType.get(shape, f32, memory_space=mgpu.utils.smem())
       lhs_ty = ref_ty if lhs_memory_space == "shared" else vec_ty
       acc, lhs, rhs = undefs(vec_ty, lhs_ty, ref_ty)
       wgmma_op = mgpu.dialect.WGMMAOp(acc, lhs, rhs)