[Pallas/TPU] Add supports_sparse_core property to TpuInfo.

PiperOrigin-RevId: 823586333

Author: Google-ML-Automation

jax/_src/pallas/mosaic/tpu_info.py

@@ -78,6 +78,12 @@ def is_split_chip(self) -> bool:
     # Is this a multi-core chip being used in single-core mode?
     return self.num_cores == 1 and not self.is_lite
 
+  @property
+  def supports_sparse_core(self) -> bool:
+    return (
+        self.chip_version == ChipVersion.TPU_V5P and not self.is_split_chip
+    ) or self.chip_version in {ChipVersion.TPU_7X, ChipVersion.TPU_V6E}
+
   def is_matmul_supported(
       self,
       lhs_dtype: jnp.dtype | str,