[pallas:sc] Skipped a few more tests under TC tiling due to verification errors

PiperOrigin-RevId: 833257620

Author: superbobry

jaxlib/mosaic/dialect/tpu/tpu_dialect.h

@@ -59,7 +59,8 @@ std::pair<bool, bool> mightCommunicateBetweenChips(Operation *op);
 
 std::unique_ptr<OperationPass<func::FuncOp>> createInferMemRefLayoutPass(
     int hardware_generation, std::array<int64_t, 2> target_shape,
-    const TpuTilingFlags& tpu_tiling_flags, bool align = true);
+    const TpuTilingFlags& tpu_tiling_flags, bool align = true,
+    bool infer_kernel_arguments = true);
 
 #define GEN_PASS_DECL_MOSAICSERDEPASS
 #include "jaxlib/mosaic/dialect/tpu/tpu_passes.h.inc"

tests/pallas/tpu_sparsecore_pallas_test.py

@@ -296,7 +296,13 @@ def kernel(x_ref, o_ref):
   @hp.given(hps.data())
   def test_block_spec_untiled_slicing(self, data):
     if not self.USE_TC_TILING:
-      self.skipTest("Test uncoveres a bug: @reproduce_failure('6.80.0', b'AAEBAQAAAAA=')")
+      self.skipTest(
+          "Test uncovers a bug: @reproduce_failure('6.80.0', b'AAEBAQAAAAA=')"
+      )
+    else:
+      self.skipTest(
+          "Test uncovers a bug: @reproduce_failure('6.80.0', b'AAEAAQAAAAA=')"
+      )
     slice_shape = data.draw(
         hps.lists(
             hps.integers(1, 3), min_size=(1 + self.USE_TC_TILING), max_size=4
@@ -1065,7 +1071,7 @@ def scoped_kernel(scratch_ref):
       pl.run_scoped(
           scoped_kernel,
           plsc.MemoryRef(
-              x.shape, x_ref.dtype, memory_space=pltpu.VMEM, tiling=[(2, 1)]
+              x.shape, x_ref.dtype, memory_space=pltpu.VMEM, tiling=[(1, 8)]
           ),
       )