Skip to content

Commit 0772c77

Browse files
superbobryGoogle-ML-Automation
authored andcommitted
[pallas:mosaic] Inlined tpu_info.supports_sparse_core
Users can check whether SparseCore is available by inspecting `tpu_info.sparse_core`. PiperOrigin-RevId: 825920709
1 parent 415e819 commit 0772c77

File tree

1 file changed

+0
-7
lines changed

1 file changed

+0
-7
lines changed

jax/_src/pallas/mosaic/tpu_info.py

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -87,13 +87,6 @@ def is_split_chip(self) -> bool:
8787
# Is this a multi-core chip being used in single-core mode?
8888
return self.num_cores == 1 and not self.is_lite
8989

90-
@property
91-
def supports_sparse_core(self) -> bool:
92-
# TODO(slebedev): Migrate the callers to use ``sparse_core`` directly.
93-
return (
94-
self.chip_version == ChipVersion.TPU_V5P and not self.is_split_chip
95-
) or self.chip_version in {ChipVersion.TPU_7X, ChipVersion.TPU_V6E}
96-
9790
def is_matmul_supported(
9891
self,
9992
lhs_dtype: jnp.dtype | str,

0 commit comments

Comments
 (0)