[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

@@ -297,6 +297,10 @@ def kernel(x_ref, o_ref):
   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=')")
+    else:
+      self.skipTest(
+          "Test uncoveres 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
@@ -438,6 +442,8 @@ def kernel(x_hbm_ref, indices_ref, o_ref):
     np.testing.assert_array_equal(kernel(x, indices)[0], x[indices])
 
   def test_large_gather_1d(self):
+    self.skip_if_tc_tiling()
+
     x = jnp.arange(1024)
     indices = jax.random.permutation(jax.random.key(42), x)
 
@@ -471,6 +477,8 @@ def kernel(x_hbm_ref, indices_ref, o_ref):
     np.testing.assert_array_equal(kernel(x, indices), x[1, 2, indices])
 
   def test_gather_2d_with_indexing(self):
+    self.skip_if_tc_tiling()
+
     x = jnp.arange(4 * 16 * 128).reshape(4, 16, 128)
     indices = jax.random.permutation(jax.random.key(42), jnp.arange(8))
 
@@ -997,6 +1005,7 @@ def kernel(x_ref, o_ref):
     np.testing.assert_array_equal(kernel(x)[5:13:2], x[2:6])
 
   def test_scalar_load_store(self):
+    self.skip_if_tc_tiling()
 
     @self.vector_subcore_kernel(
         in_specs=(pl.BlockSpec(memory_space=pltpu.HBM),),
@@ -1054,6 +1063,8 @@ def kernel(x_ref, o_hbm_ref):
     np.testing.assert_array_equal(kernel(x), x)
 
   def test_run_scoped_with_tiling(self):
+    self.skip_if_tc_tiling()
+
     x = jnp.arange(2 * 8).reshape(-1, 8)
 
     @self.vector_subcore_kernel(out_shape=x)
@@ -1089,6 +1100,8 @@ def kernel(x_ref, o_ref):
     np.testing.assert_array_equal(kernel(x), x)
 
   def test_scratch(self):
+    self.skip_if_tc_tiling()
+
     x = jnp.arange(8)
 
     @self.vector_subcore_kernel(