Add: Add pathwayutils to core in pyproject.toml to bypass modul import error

Author: camiloCienet

axlearn/common/test_utils.py

@@ -15,12 +15,6 @@
 from typing import Any, NamedTuple, Optional, Protocol, TypeVar, Union
 from unittest.mock import patch
 
-# For bypass pathwaysutil import error
-import sys
-from unittest.mock import MagicMock
-fakePathwaysuitls = MagicMock()
-sys.modules['pathwaysutils'] = fakePathwaysuitls
-
 import jax
 import jax.random
 import numpy as np

pyproject.toml

@@ -33,6 +33,7 @@ core = [
     "optax==0.2.6",  # Fix for TPU 7x
     "portpicker",
     "pyarrow<21.0.0",  # Pin to v20.x to avoid PyExtensionType -> ExtensionType breaking change in v21
+    "pathwaysutils @ git+https://github.com/AI-Hypercomputer/pathways-utils.git",
     "protobuf>4",
     "tensorboard-plugin-profile==2.20.4",
     # This has both x86 and arm64 wheels. Underneath the hood it uses tensorflow-macos since 2.13.