From a4ad9875c8c2c2934fe6b0da595d2b92f2d3f501 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Mon, 2 Nov 2026 23:08:59 -0800 Subject: [PATCH] Add `active_devices` property to `pathwaysutils.elastic.manager.Manager`. PiperOrigin-RevId: 1012738354 --- pathwaysutils/elastic/manager.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pathwaysutils/elastic/manager.py b/pathwaysutils/elastic/manager.py index 8f4085a..a6d2bc5 100644 --- a/pathwaysutils/elastic/manager.py +++ b/pathwaysutils/elastic/manager.py @@ -198,6 +198,14 @@ def active_slice_count(self) -> int: """The number of active slices.""" return len(self.active_slice_indices) + @property + def active_devices(self) -> list[jax.Device]: + """The devices belonging to the active slices, ordered by slice index.""" + devices = [] + for slice_index in sorted(self.active_slice_indices): + devices.extend(self.slice_to_devices[slice_index]) + return devices + def scale_by_active_slices(self, x: int | float) -> int | float: """Scale x by the number of active slices.""" if isinstance(x, int):