Fix: New calculation for num_granule.

Author: camiloCienet

axlearn/common/utils.py

@@ -1830,9 +1830,9 @@ def create_device_mesh(
     if not all(el.platform == device_platform for el in devices):
         raise NotImplementedError(f"Not all devices had platform: {device_platform}.")
     print("IS MULTI GRANULE by camilo: {is_multi_granule_env}")
-    print(f"Granule devices: {(getattr(el, device_attr) for el in devices.flatten())} + 1, by camilo")
+    print(f"Granule devices: {[getattr(el, device_attr) for el in devices.flatten()]} + 1, by camilo")
     num_granules = (
-        max(getattr(el, device_attr) for el in devices.flatten()) + 1 if is_multi_granule_env else 1
+        len({getattr(el, device_attr) for el in devices.flatten()})  if is_multi_granule_env else 1
     )
     num_devices = len(devices)
     assert (

pyproject.toml

@@ -109,7 +109,7 @@ gcp = [
 tpu = [
     "axlearn[gcp]",
     "jax[tpu]==0.10.0",  # must be >=0.4.19 for compat with v5p.
-    "pathwaysutils==0.1.2",  # For JAX+Pathways single-controller accelerator coordinator.
+    "pathwaysutils @ git+https://github.com/AI-Hypercomputer/pathways-utils.git",
 ]
 # Vertex AI tensorboard. TODO(markblee): Merge with `gcp`.
 vertexai_tensorboard = [
@@ -174,6 +174,9 @@ build.targets.wheel.exclude = ["**/BUILD", "axlearn/experiments/testdata"]
 # This is where the code is for historical reasons.
 build.targets.wheel.packages = ["axlearn"]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [tool.pytype]
 inputs = ["axlearn"]