Simple fix

pyproject.toml

@@ -66,7 +66,6 @@ only-include = ["pyproject.toml", "README.md"]
 [tool.basedpyright] 
 typeCheckingMode = "strict"
 failOnWarnings = true 
-stubPath = "stubs"
 
 reportAny = "error"
 reportUnknownVariableType = "error"

shared/protobufs/types/mlx/nn/__init__.pyi

@@ -0,0 +1,3 @@
+from mlx.nn.layers import *
+from mlx.nn import init as init, losses as losses
+from mlx.nn.utils import average_gradients as average_gradients, value_and_grad as value_and_grad

shared/types/worker/commands_runner.py

@@ -11,7 +11,7 @@ from shared.openai import FinishReason
 
 class MessageType(str, Enum):
     Setup = 'setup'
-    ChatTask = "chat_task"
+    ChatTask = 'chat_task'
     Exit = 'exit'
 
 MT = TypeVar(name='MT', bound=MessageType)

shared/types/worker/resource_monitor.py

@@ -3,11 +3,11 @@ from collections.abc import Coroutine
 
 import asyncio
 
-from shared.types.events.events import ResourceProfiledEvent
+from shared.types.events.events import ResourceProfiled
 from shared.types.profiling.common import NodePerformanceProfile, MemoryPerformanceProfile, SystemPerformanceProfile
 
 class EventLog:
-    def append(self, event: ResourceProfiledEvent) -> None:
+    def append(self, event: ResourceProfiled) -> None:
         ...
 
 class ResourceCollector(ABC):

shared/utils.py

@@ -2,7 +2,7 @@ from typing import Any, Type, TypeVar
 
 T = TypeVar('T')
 
-def ensure_type(obj: Any, expected_type: Type[T]) -> T:
+def ensure_type(obj: Any, expected_type: Type[T]) -> T: # type: ignore
     if not isinstance(obj, expected_type):
-        raise TypeError(f"Expected {expected_type}, got {type(obj)}")
+        raise TypeError(f"Expected {expected_type}, got {type(obj)}") # type: ignore
     return obj