From 374e726fdc64d04d76e1608a65f08fb76807941c Mon Sep 17 00:00:00 2001 From: Konrad Kaim Date: Tue, 21 Oct 2025 16:19:52 +0000 Subject: [PATCH] feat: use get_topology_product instead of calculating topology product manually --- src/xpk/core/system_characteristics.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/xpk/core/system_characteristics.py b/src/xpk/core/system_characteristics.py index 62d6ce3b8..b0aac13ce 100644 --- a/src/xpk/core/system_characteristics.py +++ b/src/xpk/core/system_characteristics.py @@ -15,8 +15,8 @@ """ from dataclasses import dataclass -from functools import reduce -from operator import mul +from ..utils.topology import get_topology_product + AcceleratorType = {'TPU': 1, 'GPU': 2, 'CPU': 3} @@ -136,7 +136,7 @@ def get_tpu_system_characteristics_map( ) -> dict[str, SystemCharacteristics]: system_characteristics_map = {} for topology in supported_topologies: - total_chips = reduce(mul, (int(x) for x in topology.split('x')), 1) + total_chips = get_topology_product(topology) num_tensorcores = total_chips * tensorcores_per_chip chips_per_vm = 1 if total_chips == 1 else 4 vms_per_slice = total_chips // chips_per_vm