remove cbrt which doesnt exist on python 3.9

4a7503551a4231c8ca869aed54b112e6c13bd498 · 2024-10-23 14:46:41 -0700 · Alex Cheema

Files touched

M exo/viz/topology_viz.py

Diff

commit 4a7503551a4231c8ca869aed54b112e6c13bd498
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Oct 23 14:46:41 2024 -0700

    remove cbrt which doesnt exist on python 3.9
---
 exo/viz/topology_viz.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/viz/topology_viz.py b/exo/viz/topology_viz.py
index ab66773a..b9910e9d 100644
--- a/exo/viz/topology_viz.py
+++ b/exo/viz/topology_viz.py
@@ -161,7 +161,7 @@ class TopologyViz:
 
     # Calculate total FLOPS and position on the bar
     total_flops = sum(self.topology.nodes.get(partition.node_id, UNKNOWN_DEVICE_CAPABILITIES).flops.fp16 for partition in self.partitions)
-    bar_pos = (math.tanh(math.cbrt(total_flops)/2.5 - 2) + 1)
+    bar_pos = (math.tanh(total_flops**(1/3)/2.5 - 2) + 1)
 
     # Add GPU poor/rich bar
     bar_width = 30

← 92270595 removed logging · back to Exo · simplify manual discovery 6e87d7c0 →