diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 74681438..9dcf41c4 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -681,8 +681,8 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
         def __call__(
             self,
             x: mx.array,
-            mask: Optional[mx.array] = None,
-            cache: Optional[Any] = None,
+            mask: mx.array | Any = None,
+            cache: Any | None = None,
         ) -> mx.array:
             B, L, D = x.shape