Enable gather_with_index test on tc tiling. Also add a slicing to the test.

PiperOrigin-RevId: 834957962

Author: IvyZX

tests/pallas/tpu_sparsecore_pallas_test.py

@@ -416,22 +416,21 @@ def kernel(x_hbm_ref, indices_ref, o_ref):
 
     np.testing.assert_array_equal(kernel(x, indices), x[indices])
 
-  def test_gather_1d_with_indexing(self):
-    self.skip_if_tc_tiling()
-    x = jnp.arange(4 * 4 * 8).reshape(4, 4, 8)
+  def test_gather_2d_with_indexing(self):
+    x = jnp.arange(4 * 16 * 128).reshape(4, 16, 128)
     indices = jax.random.permutation(jax.random.key(42), jnp.arange(8))
 
     @self.vector_subcore_kernel(
-        out_shape=jax.ShapeDtypeStruct(shape=(8,), dtype=jnp.int32),
+        out_shape=jax.ShapeDtypeStruct(shape=(8, 128,), dtype=jnp.int32),
         in_specs=(
             pl.BlockSpec(memory_space=pltpu.HBM),
             pl.BlockSpec(memory_space=pltpu.VMEM),
         ),
     )
     def kernel(x_hbm_ref, indices_ref, o_ref):
-      pltpu.sync_copy(x_hbm_ref.at[1, 2].at[indices_ref], o_ref)
+      pltpu.sync_copy(x_hbm_ref.at[1, pl.ds(8, 8), :].at[indices_ref], o_ref)
 
-    np.testing.assert_array_equal(kernel(x, indices), x[1, 2, indices])
+    np.testing.assert_array_equal(kernel(x, indices), x[1, 8:][indices])
 
   def test_gather_1d_with_indexed_ref(self):
     x = jnp.arange(16)